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

[aartbik]

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

803–804

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 ↗

(On Diff #476195)

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 ↗

(On Diff #476195)

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/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();
+  Value input = op.getTensor();
+  SmallVector<Value> reduc = op.getInitArgs();
+  auto rtp = input.getType().cast<RankedTensorType>();
+  int64_t rank = rtp.getRank();
+
+  // Foreach on constant.
+  DenseElementsAttr indicesAttr = attr.getIndices();
+  DenseElementsAttr valuesAttr = attr.getValues();
+
+  SmallVector<Value> args;
+  for (int i = 0, e = valuesAttr.size(); i < e; 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);
+    }
+    Value val;
+    if (rtp.getElementType().isa<ComplexType>()) {
+      auto valAttr = valuesAttr.getValues<ArrayAttr>()[i];
+      val = rewriter.create<complex::ConstantOp>(loc, rtp.getElementType(),
+                                                 valAttr);
+    } else {
+      auto valAttr = valuesAttr.getValues<TypedAttr>()[i];
+      // Remaps value.
+      val = rewriter.create<arith::ConstantOp>(loc, valAttr);
+    }
+    assert(val);
+    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();
+}
+
 //===---------------------------------------------------------------------===//
 // 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)
   ],