[mlir][sparse] fix crash on sparse_tensor.foreach operation on tensors with complex<T> elements.

[aartbik]

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

783–784

Can we push this into a util?

[Peiming]

move foreach on constant in a separate util function.

[Peiming]

go deeper

[wrengr]

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

1047

Should we also be checking for f16/bf16? I ask because those two generally pattern together with complex when it comes to issues like this

[Peiming]

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

1047

I assume that ArrayAttr carrys the f16 will be typed, so that the complexOp will be set to the right element type.

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

 
 /// Populates given sizes array from dense tensor or sparse tensor constant.
 void sizesFromSrc(OpBuilder &builder, SmallVectorImpl<Value> &sizes,
                   Location loc, Value src);
 
 /// Scans to top of generated loop.
 Operation *getTop(Operation *op);
 
+/// Iterate over a sparse constant, generates constantOp for value and indices.
+/// E.g.,
+/// sparse<[ [0], [28], [31] ],
+///          [ (-5.13, 2.0), (3.0, 4.0), (5.0, 6.0) ] >
+/// =>
+/// %c1 = arith.constant 0
+/// %v1 = complex.constant (5.13, 2.0)
+/// callback({%c1}, %v1)
+///
+/// %c2 = arith.constant 28
+/// %v2 = complex.constant (3.0, 4.0)
+/// callback({%c2}, %v2)
+///
+/// %c3 = arith.constant 31
+/// %v3 = complex.constant (5.0, 6.0)
+/// callback({%c3}, %v3)
+void foreachInSparseConstant(
+    Location loc, RewriterBase &rewriter, SparseElementsAttr attr,
+    function_ref<void(ArrayRef<Value>, Value)> callback);
+
 //===----------------------------------------------------------------------===//
 // Inlined constant generators.
 //
 // All these functions are just wrappers to improve code legibility;
 // therefore, we mark them as `inline` to avoid introducing any additional
 // overhead due to the legibility.
 //
 // TODO: Ideally these should move upstream, so that we don't
 // develop a design island.  However, doing so will involve
 // substantial design work.  For related prior discussion, see
 // <https://llvm.discourse.group/t/evolving-builder-apis-based-on-lessons-learned-from-edsc/879>
 //===----------------------------------------------------------------------===//
 
 /// Generates a 0-valued constant of the given type.  In addition to
-/// the scalar types (`ComplexType`, ``FloatType`, `IndexType`, `IntegerType`),
-/// this also works for `RankedTensorType` and `VectorType` (for which it
-/// generates a constant `DenseElementsAttr` of zeros).
+/// the scalar types (`ComplexType`, ``FloatType`, `IndexType`,
+/// `IntegerType`), this also works for `RankedTensorType` and `VectorType`
+/// (for which it generates a constant `DenseElementsAttr` of zeros).
 inline Value constantZero(OpBuilder &builder, Location loc, Type tp) {
   if (auto ctp = tp.dyn_cast<ComplexType>()) {
     auto zeroe = builder.getZeroAttr(ctp.getElementType());
     auto zeroa = builder.getArrayAttr({zeroe, zeroe});
     return builder.create<complex::ConstantOp>(loc, tp, zeroa);
   }
   return builder.create<arith::ConstantOp>(loc, tp, builder.getZeroAttr(tp));
 }

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

Operation *mlir::sparse_tensor::getTop(Operation *op) {
   for (; isa<scf::ForOp>(op->getParentOp()) ||
          isa<scf::WhileOp>(op->getParentOp()) ||
          isa<scf::ParallelOp>(op->getParentOp()) ||
          isa<scf::IfOp>(op->getParentOp());
        op = op->getParentOp())
     ;
   return op;
 }
+
+void sparse_tensor::foreachInSparseConstant(
+    Location loc, RewriterBase &rewriter, SparseElementsAttr attr,
+    function_ref<void(ArrayRef<Value>, Value)> callback) {
+  int64_t rank = attr.getType().getRank();
+  // Foreach on constant.
+  DenseElementsAttr indicesAttr = attr.getIndices();
+  DenseElementsAttr valuesAttr = attr.getValues();
+
+  SmallVector<Value> coords;
+  for (int i = 0, e = valuesAttr.size(); i < e; i++) {
+    coords.clear();
+    for (int j = 0; j < rank; j++) {
+      auto coordAttr = indicesAttr.getValues<IntegerAttr>()[i * rank + j];
+      auto coord =
+          rewriter.create<arith::ConstantIndexOp>(loc, coordAttr.getInt());
+      // Remaps coordinates.
+      coords.push_back(coord);
+    }
+    Value val;
+    if (attr.getElementType().isa<ComplexType>()) {
+      auto valAttr = valuesAttr.getValues<ArrayAttr>()[i];
+      val = rewriter.create<complex::ConstantOp>(loc, attr.getElementType(),
+                                                 valAttr);
+    } else {
+      auto valAttr = valuesAttr.getValues<TypedAttr>()[i];
+      // Remaps value.
+      val = rewriter.create<arith::ConstantOp>(loc, valAttr);
+    }
+    assert(val);
+    callback(coords, val);
+  }
+}

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

static void getDynamicSizes(RankedTensorType tp,
                             const SmallVectorImpl<Value> &sizes,
                             SmallVectorImpl<Value> &dynSizes) {
   for (const auto &d : enumerate(tp.getShape())) {
     if (d.value() == ShapedType::kDynamicSize)
       dynSizes.push_back(sizes[d.index()]);
   }
 }
 
+static LogicalResult genForeachOnSparseConstant(ForeachOp op,
+                                                RewriterBase &rewriter,
+                                                SparseElementsAttr attr) {
+  auto loc = op.getLoc();
+  SmallVector<Value> reduc = op.getInitArgs();
+
+  // Foreach on constant.
+  foreachInSparseConstant(
+      loc, rewriter, attr,
+      [&reduc, &rewriter, op](ArrayRef<Value> coords, Value v) mutable {
+        SmallVector<Value> args;
+        args.append(coords.begin(), coords.end());
+        args.push_back(v);
+        args.append(reduc);
+        // Clones the foreach op to get a copy of the loop body.
+        auto cloned = cast<ForeachOp>(rewriter.clone(*op.getOperation()));
+        assert(args.size() == cloned.getBody()->getNumArguments());
+        Operation *yield = cloned.getBody()->getTerminator();
+        rewriter.mergeBlockBefore(cloned.getBody(), op, args);
+        // clean up
+        rewriter.eraseOp(cloned);
+        reduc = yield->getOperands();
+        rewriter.eraseOp(yield);
+      });
+
+  rewriter.replaceOp(op, reduc);
+  return success();
+}
+
 //===---------------------------------------------------------------------===//
 // The actual sparse tensor rewriting rules.
 //===---------------------------------------------------------------------===//
 
 namespace {
 
 /// Rewriting rule that converts direct yield of zero with initial allocation.
 struct FoldInvariantYield : public OpRewritePattern<GenericOp> {
LogicalResult matchAndRewrite(ForeachOp op,
     Value input = op.getTensor();
     SmallVector<Value> reduc = op.getInitArgs();
     auto rtp = input.getType().cast<RankedTensorType>();
     int64_t rank = rtp.getRank();
 
     // Special-case: for each over a sparse constant uses its own rewriting
     // rule.
     if (auto constOp = input.getDefiningOp<arith::ConstantOp>()) {
       if (auto attr = constOp.getValue().dyn_cast<SparseElementsAttr>()) {
-        // Foreach on constant.
-        DenseElementsAttr indicesAttr = attr.getIndices();
-        DenseElementsAttr valuesAttr = attr.getValues();
-
-        SmallVector<Value> args;
-        for (int i = 0, e = valuesAttr.size(); i < e; i++) {
-          auto valAttr = valuesAttr.getValues<TypedAttr>()[i];
-          for (int j = 0; j < rank; j++) {
-            auto coordAttr = indicesAttr.getValues<IntegerAttr>()[i * rank + j];
-            auto coord = rewriter.create<arith::ConstantIndexOp>(
-                loc, coordAttr.getInt());
-            // Remaps coordinates.
-            args.push_back(coord);
-          }
-          // Remaps value.
-          auto val = rewriter.create<arith::ConstantOp>(loc, valAttr);
-          args.push_back(val);
-          // Remaps iteration args.
-          args.append(reduc);
-          auto cloned = cast<ForeachOp>(rewriter.clone(*op.getOperation()));
-          Operation *yield = cloned.getBody()->getTerminator();
-          rewriter.mergeBlockBefore(cloned.getBody(), op, args);
-          // clean up
-          args.clear();
-          rewriter.eraseOp(cloned);
-          reduc = yield->getOperands();
-          rewriter.eraseOp(yield);
-        }
-        rewriter.replaceOp(op, reduc);
-        return success();
+        return genForeachOnSparseConstant(op, rewriter, attr);
       }
     }
 
     // Otherwise, use loop emitter to generate loops.
     auto enc = getSparseTensorEncoding(rtp);
 
     // 1. Generates loop for the sparse input.
     SparseTensorLoopEmitter loopEmitter(

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_complex_ops.mlir

-// RUN: mlir-opt %s --sparse-compiler | \
-// RUN: mlir-cpu-runner \
-// RUN:  -e entry -entry-point-result=void  \
-// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
-// RUN: FileCheck %s
+// DEFINE: %{option} = enable-runtime-library=true
+// DEFINE: %{command} = mlir-opt %s --sparse-compiler=%{option} | \
+// DEFINE: mlir-cpu-runner \
+// DEFINE:  -e entry -entry-point-result=void  \
+// DEFINE:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// DEFINE: FileCheck %s
+//
+// RUN: %{command}
+//
+// Do the same run, but now with direct IR generation.
+// REDEFINE: %{option} = enable-runtime-library=false
+// RUN: %{command}
 
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
 
 #trait_op1 = {
   indexing_maps = [
     affine_map<(i) -> (i)>,  // a (in)
     affine_map<(i) -> (i)>   // x (out)
   ],