[mlir][sparse] Refactor the conversion of the tensor reshape operators.

[bixia]

Summary

Move genReshapeDstShape to codegen utils to support the rewriting of the tensor
reshape operators for the codegen path.

[bixia]

Fix a comment.

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h

 /// Generates a constant of the internal dimension level type encoding.
 inline Value
 constantDimLevelTypeEncoding(OpBuilder &builder, Location loc,
                              SparseTensorEncodingAttr::DimLevelType dlt) {
   return constantI8(builder, loc,
                     static_cast<uint8_t>(dimLevelTypeEncoding(dlt)));
 }
 
+/// Compute the shape of destination tensor of a reshape operator. This is only
+/// used when operands have dynamic shape. The shape of the destination is
+/// stored into dstShape.
+void genReshapeDstShape(Location loc, PatternRewriter &rewriter,
+                        SmallVector<Value, 4> &dstShape,
+                        ArrayRef<Value> srcShape,
+                        ArrayRef<int64_t> staticDstShape,
+                        ArrayRef<ReassociationIndices> reassociation);
+
 /// Helper method to translate indices during a reshaping operation.
 void translateIndicesArray(OpBuilder &builder, Location loc,
                            ArrayRef<ReassociationIndices> reassociation,
                            ValueRange srcIndices, ArrayRef<Value> srcShape,
                            ArrayRef<Value> dstShape,
                            SmallVectorImpl<Value> &dstIndices);
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENUTILS_H_

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp

Value mlir::sparse_tensor::genIsNonzero(OpBuilder &builder, mlir::Location loc,
   if (tp.isIntOrIndex())
     return builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ne, v,
                                          zero);
   if (tp.dyn_cast<ComplexType>())
     return builder.create<complex::NotEqualOp>(loc, v, zero);
   llvm_unreachable("Non-numeric type");
 }
 
+void mlir::sparse_tensor::genReshapeDstShape(
+    Location loc, PatternRewriter &rewriter, SmallVector<Value, 4> &dstShape,
+    ArrayRef<Value> srcShape, ArrayRef<int64_t> staticDstShape,
+    ArrayRef<ReassociationIndices> reassociation) {
+  // Collapse shape.
+  if (reassociation.size() < srcShape.size()) {
+    unsigned start = 0;
+    for (const auto &map : llvm::enumerate(reassociation)) {
+      auto dstDim = constantIndex(rewriter, loc, 1);
+      for (unsigned i = start; i < start + map.value().size(); i++) {
+        dstDim = rewriter.create<arith::MulIOp>(loc, dstDim, srcShape[i]);
+      }
+      dstShape.push_back(dstDim);
+      start = start + map.value().size();
+    }
+    assert(start == srcShape.size());
+    return;
+  }
+
+  // Expand shape.
+  assert(reassociation.size() == srcShape.size());
+  unsigned start = 0;
+  // Expand the i-th dimension in srcShape.
+  for (unsigned i = 0, size = srcShape.size(); i < size; i++) {
+    auto map = reassociation[i];
+    auto srcDim = srcShape[i];
+    // Iterate through dimensions expanded from the i-th dimension.
+    for (unsigned j = start; j < start + map.size(); j++) {
+      // There can be only one dynamic sized dimension among dimensions expanded
+      // from the i-th dimension in srcShape. For example, if srcDim = 8, then
+      // the expanded shape could be <2x?x2>, but not <2x?x?>.
+      if (staticDstShape[j] == ShapedType::kDynamicSize) {
+        // The expanded dimension has dynamic size. We compute the dimension
+        // by dividing srcDim by the product of the static dimensions.
+        int64_t product = 1;
+        for (unsigned k = start; k < start + map.size(); k++) {
+          if (staticDstShape[k] != ShapedType::kDynamicSize) {
+            product *= staticDstShape[k];
+          }
+        }
+        // Compute the dynamic dimension size.
+        Value productVal = constantIndex(rewriter, loc, product);
+        Value dynamicSize =
+            rewriter.create<arith::DivUIOp>(loc, srcDim, productVal);
+        dstShape.push_back(dynamicSize);
+      } else {
+        // The expanded dimension is statically known.
+        dstShape.push_back(constantIndex(rewriter, loc, staticDstShape[j]));
+      }
+    }
+    start = start + map.size();
+  }
+  assert(start == staticDstShape.size());
+}
+
 void mlir::sparse_tensor::translateIndicesArray(
     OpBuilder &builder, Location loc,
     ArrayRef<ReassociationIndices> reassociation, ValueRange srcIndices,
     ArrayRef<Value> srcShape, ArrayRef<Value> dstShape,
     SmallVectorImpl<Value> &dstIndices) {
   unsigned i = 0;
   unsigned start = 0;
   unsigned dstRank = dstShape.size();

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp

static void translateIndices(Location loc, ConversionPatternRewriter &rewriter,
   translateIndicesArray(rewriter, loc, reassociation, srcIndices, srcShape,
                         dstShape, dstIndices);
 
   for (unsigned i = 0; i < dstRank; i++)
     rewriter.create<memref::StoreOp>(loc, dstIndices[i], dstIdx,
                                      constantIndex(rewriter, loc, i));
 }
 
-/// Helper method to compute the shape of destination tensor of a reshape
-/// operator. This is only used when operands have dynamic shape. The shape of
-/// the destination is stored into dstShape.
-void genReshapeDstShape(Location loc, ConversionPatternRewriter &rewriter,
-                        SmallVector<Value, 4> &dstShape,
-                        ArrayRef<Value> srcShape,
-                        ArrayRef<int64_t> staticDstShape,
-                        ArrayRef<ReassociationIndices> reassociation) {
-  // Collapse shape.
-  if (reassociation.size() < srcShape.size()) {
-    unsigned start = 0;
-    for (const auto &map : llvm::enumerate(reassociation)) {
-      auto dstDim = constantIndex(rewriter, loc, 1);
-      for (unsigned i = start; i < start + map.value().size(); i++) {
-        dstDim = rewriter.create<arith::MulIOp>(loc, dstDim, srcShape[i]);
-      }
-      dstShape.push_back(dstDim);
-      start = start + map.value().size();
-    }
-    assert(start == srcShape.size());
-    return;
-  }
-
-  // Expand shape.
-  assert(reassociation.size() == srcShape.size());
-  unsigned start = 0;
-  // Expand the i-th dimension in srcShape.
-  for (unsigned i = 0, size = srcShape.size(); i < size; i++) {
-    auto map = reassociation[i];
-    auto srcDim = srcShape[i];
-    // Iterate through dimensions expanded from the i-th dimension.
-    for (unsigned j = start; j < start + map.size(); j++) {
-      // There can be only one dynamic sized dimension among dimensions expanded
-      // from the i-th dimension in srcShape. For example, if srcDim = 8, then
-      // the expanded shape could be <2x?x2>, but not <2x?x?>.
-      if (staticDstShape[j] == ShapedType::kDynamicSize) {
-        // The expanded dimension has dynamic size. We compute the dimension
-        // by dividing srcDim by the product of the static dimensions.
-        int64_t product = 1;
-        for (unsigned k = start; k < start + map.size(); k++) {
-          if (staticDstShape[k] != ShapedType::kDynamicSize) {
-            product *= staticDstShape[k];
-          }
-        }
-        // Compute the dynamic dimension size.
-        Value productVal = constantIndex(rewriter, loc, product);
-        Value dynamicSize =
-            rewriter.create<arith::DivUIOp>(loc, srcDim, productVal);
-        dstShape.push_back(dynamicSize);
-      } else {
-        // The expanded dimension is statically known.
-        dstShape.push_back(constantIndex(rewriter, loc, staticDstShape[j]));
-      }
-    }
-    start = start + map.size();
-  }
-  assert(start == staticDstShape.size());
-}
-
 /// Generate code for a general sparse to sparse reshaping operation.
 /// Note that unlike dense reshaping (which can be done with a "cheap"
 /// change of view), sparse reshaping is currently done with actual
 /// data shuffling.
 ///
 /// TODO: proportional to nnz, but still a lot of data movement
 ///       https://github.com/llvm/llvm-project/issues/56477
 ///