[mlir][sparse] extend foreach operation to iterator over sparse constant.

[aartbik]

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

703

I would add a comment before L703

// Special-case: for each over a sparse constant uses its own rewriting rule.

Just to set it apart from the 1.2. steps below

also, perhaps later, perhaps we should move some of these larger codeblocks into their own helper method

[Peiming]

address comments.

[Peiming]

rebase

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

 public:
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ForeachOp op,
                                 PatternRewriter &rewriter) const override {
 
     auto loc = op.getLoc();
     Value input = op.getTensor();
+    SmallVector<Value> reduc = op.getInitArgs();
     auto rtp = input.getType().cast<RankedTensorType>();
     int64_t rank = rtp.getRank();
-    auto enc = getSparseTensorEncoding(rtp);
 
-    SmallVector<Value> reduc = op.getInitArgs();
+    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();
+      }
+    }
+
+    auto enc = getSparseTensorEncoding(rtp);
 
     // 1. Generates loop for the sparse input.
     SparseTensorLoopEmitter loopEmitter(ValueRange{input});
     loopEmitter.initializeLoopEmit(rewriter, loc);
     for (int64_t i = 0; i < rank; i++) {
       // TODO: provide utility function for loop sequences that only contains
       // one for loop?
       loopEmitter.enterNewLoopSeq(rewriter, loc, 0, static_cast<size_t>(i));

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

 }>
 
 #SortedCOOPerm = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed-nu", "singleton" ],
   dimOrdering = affine_map<(i,j) -> (j,i)>
 }>
 
 module {
+  /// uses foreach operator to print coords and values.
+  func.func @foreach_print_const() {
+    // Initialize a tensor.
+    %0 = arith.constant sparse<[[0, 0], [1, 6]], [1.0, 5.0]> : tensor<8x7xf32>
+    sparse_tensor.foreach in %0 : tensor<8x7xf32> do {
+      ^bb0(%1: index, %2: index, %v: f32) :
+        vector.print %1: index
+        vector.print %2: index
+        vector.print %v: f32
+     }
+     return
+  }
 
   /// uses foreach operator to print coords and values.
   func.func @foreach_print_1(%arg0: tensor<2x2xf64, #Row>) {
     sparse_tensor.foreach in %arg0 : tensor<2x2xf64, #Row> do {
       ^bb0(%1: index, %2: index, %v: f64) :
         vector.print %1: index
         vector.print %2: index
         vector.print %v: f64
func.func @entry() {
     %s1 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #Row>
     %s2 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #CSR>
     %s3 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #DCSC>
     %s4 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #SortedCOO>
     %s5 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #SortedCOOPerm>
     // CHECK: 0
     // CHECK-NEXT: 0
     // CHECK-NEXT: 1
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 6
+    // CHECK-NEXT: 5
+    call @foreach_print_const() : () -> ()
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 1
     // CHECK-NEXT: 0
     // CHECK-NEXT: 1
     // CHECK-NEXT: 2
     // CHECK-NEXT: 1
     // CHECK-NEXT: 0
     // CHECK-NEXT: 5
     // CHECK-NEXT: 1
     // CHECK-NEXT: 1