[mlir][sparse] extend foreach operation to accept reduction arguments.

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td

873

update doc with init/results

922–923

Do we have two example sections?

Regardless, please add an example with SSA chain and yield

mlir/test/Dialect/SparseTensor/invalid.mlir

584

can you add one with a mismatch in the yield?

[aartbik]

Good to go after address my nits.

[Peiming]

address comments from Aart

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td

869

typo: variable, and probably even plural, variables

922–923

Maybe, For example:
too, to be consistent?

[Peiming]

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td

922–923

Can I remove it?

This used to be the common parts for all sparse tensor operations, but now it seems useless since it does not bring in any new information here.

[Peiming]

fix typo + address comments from Aart

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td

def SparseTensor_YieldOp : SparseTensor_Op<"yield", [Pure, Terminator]>,
   let assemblyFormat = [{
         $result attr-dict `:` type($result)
   }];
   let hasVerifier = 1;
 }
 
 def SparseTensor_ForeachOp : SparseTensor_Op<"foreach",
     [SingleBlockImplicitTerminator<"YieldOp">]>,
-    Arguments<(ins AnyTensor:$tensor)>{
+     Arguments<(ins AnyTensor:$tensor,
+                    Variadic<AnyType>:$initArgs)>,
+     Results<(outs Variadic<AnyType>:$results)> {
   let summary = "Iterates over elements in a tensor";
   let description = [{
      Iterates over stored elements in a tensor (which are typically, but not always,
      non-zero for sparse tensors) and executes the block.
 
-     For an input tensor with rank n, the block must take n + 1 arguments. The
-     first n arguments must be Index type, together indicating the current coordinates
-     of the element being visited. The last argument must have the same type as the
+     For an input tensor with rank n, the block must take n + 1 (and additional loop
+     carried variables as described below) arguments. The first n arguments must be
+     Index type, together indicating the current coordinates of the element being visited.
+     The last argument must have the same type as the
      tensor's element type, representing the actual value loaded from the input
      tensor at the given coordinates.
 
-     Note that foreach generated loop iterates over the stored elements in the storage
-     order. However, no matter what storage order is used, the indices passed to the block
-     always obey the original dimension order.
+     `sparse_tensor.foreach` can also operate on loop-carried variables and returns
+     the final values after loop termination. The initial values of the variables are
+     passed as additional SSA operands to the "sparse_tensor.foreach" following the n + 1
+     SSA values mentioned above (n coordinate and 1 value).
+
+     The region must terminate with a "sparse_tensor.yield" that passes the current
+     values of all loop-carried variables to the next iteration, or to the
+     result, if at the last iteration. The number and static types of loop-carried
+     variables may not change with iterations.
+
+     For example:
+     ```mlir
+     %c0 = arith.constant 0 : i32
+     %ret = sparse_tensor.foreach in %0 init(%c0): tensor<?x?xi32, #DCSR>, i32 -> i32 do {
+      ^bb0(%arg1: index, %arg2: index, %arg3: i32, %iter: i32):
+        %sum = arith.add %iter, %arg3
+        sparse_tensor.yield %sum
+     }
+     ```
+
+     It is important to note that foreach generated loop iterates over the stored elements
+     in the storage order. However, no matter what storage order is used, the indices passed
+     to the block always obey the original dimension order.
 
      For example:
      ```mlir
      #COL_MAJOR = #sparse_tensor.encoding<{
        dimLevelType = [ "compressed", "compressed" ],
        dimOrdering = affine_map<(i,j) -> (j,i)>
      }>
-     
      // foreach on a column-major sparse tensor
      sparse_tensor.foreach in %0 : tensor<2x3xf64, #COL_MAJOR> do {
-      ^bb0(%row: index, %col: index, %arg3: f64): 
+      ^bb0(%row: index, %col: index, %arg3: f64):
          // [%row, %col] -> [0, 0], [1, 0], [2, 0], [0, 1], [1, 1], [2, 1]
      }
 
      #ROW_MAJOR = #sparse_tensor.encoding<{
        dimLevelType = [ "compressed", "compressed" ],
      }>
 
      // foreach on a row-major sparse tensor
      sparse_tensor.foreach in %0 : tensor<2x3xf64, #ROW_MAJOR> do {
-      ^bb0(%row: index, %col: index, %arg3: f64): 
+      ^bb0(%row: index, %col: index, %arg3: f64):
          // [%row, %col] -> [0, 0], [0, 1], [1, 0], [1, 1], [2, 0], [2, 1]
      }
 
      ```
-
-     Example:
-
-     ```mlir
-     sparse_tensor.foreach in %0 : tensor<?x?xf64, #DCSR> do {
-      ^bb0(%arg1: index, %arg2: index, %arg3: f64):
-        do something...
-     }
-     ```
   }];
 
   let builders = [
-    OpBuilder<(
-      ins "Value":$tensor,
-      "function_ref<void(OpBuilder &, Location, ValueRange)>")>
+    OpBuilder<(ins "Value":$tensor,
+      "function_ref<void(OpBuilder &, Location, ValueRange, Value, ValueRange)>")>,
+    OpBuilder<(ins "Value":$tensor,
+      "ValueRange":$iterArgs,
+      "function_ref<void(OpBuilder &, Location, ValueRange, Value, ValueRange)>")>
   ];
 
-  let regions = (region AnyRegion:$region);
-  let assemblyFormat = "`in` $tensor attr-dict `:` type($tensor)  `do` $region";
+  let regions = (region SizedRegion<1>:$region);
+  let assemblyFormat = "`in` $tensor (`init``(`$initArgs^`)`)? attr-dict"
+                       "    `:` type($tensor) (`,` type($initArgs)^)?"
+                       "  (`->` type($results)^)?  `do` $region";
   let hasVerifier = 1;
 }
 
 #endif // SPARSETENSOR_OPS

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

LogicalResult CompressOp::verify() {
   RankedTensorType ttp = getTensor().getType().cast<RankedTensorType>();
   if (ttp.getRank() != 1 + static_cast<int64_t>(getIndices().size()))
     return emitOpError("incorrect number of indices");
   return success();
 }
 
 void ForeachOp::build(
     OpBuilder &builder, OperationState &result, Value tensor,
-    function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilder) {
-  build(builder, result, tensor);
+    function_ref<void(OpBuilder &, Location, ValueRange, Value, ValueRange)>
+        bodyBuilder) {
+  build(builder, result, tensor, llvm::None, bodyBuilder);
+}
+
+void ForeachOp::build(
+    OpBuilder &builder, OperationState &result, Value tensor,
+    ValueRange initArgs,
+    function_ref<void(OpBuilder &, Location, ValueRange, Value, ValueRange)>
+        bodyBuilder) {
+  build(builder, result, initArgs.getTypes(), tensor, initArgs);
+  // Builds foreach body.
   if (!bodyBuilder)
     return;
-
   auto rtp = tensor.getType().cast<RankedTensorType>();
   int64_t rank = rtp.getRank();
 
   SmallVector<Type, 4> blockArgTypes;
   // Starts with n index.
   std::fill_n(std::back_inserter(blockArgTypes), rank, builder.getIndexType());
   // Followed by one value.
   blockArgTypes.push_back(rtp.getElementType());
 
   SmallVector<Location, 4> blockArgLocs;
   std::fill_n(std::back_inserter(blockArgLocs), rank + 1, tensor.getLoc());
 
   OpBuilder::InsertionGuard guard(builder);
   auto &region = *result.regions.front();
   Block *bodyBlock =
       builder.createBlock(&region, region.end(), blockArgTypes, blockArgLocs);
-  bodyBuilder(builder, result.location, bodyBlock->getArguments());
+  bodyBuilder(builder, result.location,
+              bodyBlock->getArguments().slice(0, rank),
+              bodyBlock->getArguments()[rank],
+              bodyBlock->getArguments().drop_front(rank + 1));
 }
 
 LogicalResult ForeachOp::verify() {
   auto t = getTensor().getType().cast<RankedTensorType>();
   auto args = getBody()->getArguments();
 
-  if (static_cast<size_t>(t.getRank()) + 1 != args.size())
+  if (static_cast<size_t>(t.getRank()) + 1 + getInitArgs().size() !=
+      args.size())
     return emitError("Unmatched number of arguments in the block");
 
+  if (getNumResults() != getInitArgs().size())
+    return emitError("Mismatch in number of init arguments and results");
+
+  if (getResultTypes() != getInitArgs().getTypes())
+    return emitError("Mismatch in types of init arguments and results");
+
+  auto yield = cast<YieldOp>(getBody()->getTerminator());
+  if (yield.getNumOperands() != getNumResults() ||
+      yield.getOperands().getTypes() != getResultTypes())
+    return emitError("Mismatch in types of yield values and results");
+
   for (int64_t i = 0, e = t.getRank(); i < e; i++)
     if (args[i].getType() != IndexType::get(getContext()))
       emitError(
           llvm::formatv("Expecting Index type for argument at index {0}", i));
 
   auto elemTp = t.getElementType();
-  auto valueTp = args.back().getType();
+  auto valueTp = args[t.getRank()].getType();
   if (elemTp != valueTp)
     emitError(llvm::formatv("Unmatched element type between input tensor and "
                             "block argument, expected:{0}, got: {1}",
                             elemTp, valueTp));
   return success();
 }
 
 LogicalResult ReduceOp::verify() {

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

LogicalResult matchAndRewrite(ReshapeOp op,
     //   %tmp = bufferization.alloc_tensor : unordered COO
     //   foreach srcCoords %srcTensor
     //     insert translateIndicesArray(srcCoords), %tmp
     //   %t = sparse_tensor.cast %tmp
     RankedTensorType cooTp = getUnorderedCOOFromType(dstTp);
     auto cooBuffer =
         rewriter.create<AllocTensorOp>(loc, cooTp, dstDynSizes).getResult();
     rewriter.create<ForeachOp>(
-        loc, srcTensor, [&](OpBuilder &builder, Location loc, ValueRange args) {
+        loc, srcTensor, llvm::None,
+        [&](OpBuilder &builder, Location loc, ValueRange args, Value v,
+            ValueRange reduc) {
           SmallVector<Value, 4> srcIndices;
           SmallVector<Value, 4> dstIndices;
           for (int64_t i = 0, e = srcTp.getRank(); i < e; i++) {
             uint64_t dim = toStoredDim(encSrc, i);
             srcIndices.push_back(args[dim]);
           }
           translateIndicesArray(builder, loc, op.getReassociationIndices(),
                                 srcIndices, srcSizes, dstSizes, dstIndices);
-          builder.create<InsertOp>(loc, args.back(), cooBuffer, dstIndices);
+          builder.create<InsertOp>(loc, v, cooBuffer, dstIndices);
           builder.create<sparse_tensor::YieldOp>(loc);
         });
 
     rewriter.replaceOpWithNewOp<ConvertOp>(op, dstTp, cooBuffer);
     return success();
   }
 };
 
LogicalResult matchAndRewrite(ConcatenateOp op,
 
     Value offset = constantIndex(rewriter, loc, 0);
     for (Value input : op.getInputs()) {
       // Builds the indexing map.
 
       // Build a for op for each input tensor to append new values into the
       // output tensor.
       rewriter.create<ForeachOp>(
-          loc, input, [&](OpBuilder &builder, Location loc, ValueRange args) {
+          loc, input, llvm::None,
+          [&](OpBuilder &builder, Location loc, ValueRange args, Value v,
+              ValueRange reduc) {
             SmallVector<Value, 4> indices;
             for (int64_t i = 0; i < rank; i++) {
               uint64_t dim =
                   toStoredDim(getSparseTensorEncoding(input.getType()), i);
               Value idx = args[dim];
               if (i == static_cast<int64_t>(conDim))
                 // transform coordinates on matching dim
                 idx = builder.create<arith::AddIOp>(loc, idx, offset);
               indices.push_back(idx);
             }
-            builder.create<InsertOp>(loc, args.back(), cooBuffer, indices);
+            builder.create<InsertOp>(loc, v, cooBuffer, indices);
             builder.create<sparse_tensor::YieldOp>(loc);
           });
       // Accumulates the offset. Note that only static-shaped inputs are allowed
       // by concatenate op verifier, which saves us from computing the offset
       // dynamically.
       auto d = input.getType().cast<RankedTensorType>().getShape()[conDim];
       assert(!ShapedType::isDynamic(d));
       offset = rewriter.create<arith::AddIOp>(loc, offset,
LogicalResult sparse2DenseRewrite(ConvertOp op,
     RankedTensorType dstTp = op.getType().cast<RankedTensorType>();
     Value src = op.getSource();
     RankedTensorType srcTp = src.getType().cast<RankedTensorType>();
 
     SmallVector<Value, 4> sizes;
     sizesForTensor(rewriter, sizes, loc, srcTp, src);
     Value dst = allocDenseTensor(rewriter, loc, dstTp, sizes);
 
-    rewriter.create<ForeachOp>(
-        loc, src, [&](OpBuilder &builder, Location loc, ValueRange args) {
-          builder.create<memref::StoreOp>(loc, args.back(), dst,
-                                          args.drop_back());
+    rewriter.create<ForeachOp>(loc, src, llvm::None,
+                               [&](OpBuilder &builder, Location loc,
+                                   ValueRange args, Value v, ValueRange reduc) {
+                                 builder.create<memref::StoreOp>(loc, v, dst,
+                                                                 args);
                                  builder.create<sparse_tensor::YieldOp>(loc);
                                });
 
     rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, dstTp, dst);
     return success();
   }
 
   // Handles sparse tensor to sparse tensor conversion as follows:
   //   if src is not COO
   //       construct a COO to represent the src
if (!isUniqueCOOType(srcTp)) {
       // going through an intermediate COO, such as cases that only change
       // the overhead types.
       SmallVector<Value, 4> dynSrcSizes;
       getDynamicSizes(srcTp, srcSizes, dynSrcSizes);
       srcTp = getUnorderedCOOFromType(srcTp);
       tmpCoo =
           rewriter.create<AllocTensorOp>(loc, srcTp, dynSrcSizes).getResult();
       rewriter.create<ForeachOp>(
-          loc, src, [&](OpBuilder &builder, Location loc, ValueRange args) {
+          loc, src, llvm::None,
+          [&](OpBuilder &builder, Location loc, ValueRange args, Value v,
+              ValueRange reduc) {
             SmallVector<Value, 4> indices;
             for (int64_t i = 0, e = srcTp.getRank(); i < e; i++) {
               uint64_t dim = toStoredDim(encSrc, i);
               indices.push_back(args[dim]);
             }
-            builder.create<InsertOp>(loc, args.back(), tmpCoo, indices);
+            builder.create<InsertOp>(loc, v, tmpCoo, indices);
             builder.create<sparse_tensor::YieldOp>(loc);
           });
       src = tmpCoo;
     }
 
     // Sort the COO tensor so that its elements are ordered via increasing
     // indices for the storage ordering of the dst tensor.
     auto dynShape = {ShapedType::kDynamicSize};
LogicalResult sparse2SparseRewrite(ConvertOp op,
     // Sort the COO tensor.
     rewriter.create<SortOp>(loc, nnz, xs, ValueRange{y});
 
     // For each element in the COO tensor, insert the element to the dst tensor.
     SmallVector<Value, 4> dynDstSizes;
     getDynamicSizes(dstTp, srcSizes, dynDstSizes);
     Value dst =
         rewriter.create<AllocTensorOp>(loc, dstTp, dynDstSizes).getResult();
-    rewriter.create<ForeachOp>(
-        loc, src, [&](OpBuilder &builder, Location loc, ValueRange args) {
+    rewriter.create<ForeachOp>(loc, src, llvm::None,
+                               [&](OpBuilder &builder, Location loc,
+                                   ValueRange args, Value v, ValueRange reduc) {
                                  SmallVector<Value, 4> indices;
-          for (int64_t i = 0, e = srcTp.getRank(); i < e; i++) {
+                                 for (int64_t i = 0, e = srcTp.getRank(); i < e;
+                                      i++) {
                                    uint64_t dim = toStoredDim(encDst, i);
                                    indices.push_back(args[dim]);
                                  }
-          builder.create<InsertOp>(loc, args.back(), dst, indices);
+                                 builder.create<InsertOp>(loc, v, dst, indices);
                                  builder.create<sparse_tensor::YieldOp>(loc);
                                });
 
     // Release the temporary COO if it is created.
     if (tmpCoo)
       rewriter.create<DeallocTensorOp>(loc, tmpCoo);
 
     // Directly replace op with dst results in bufferization error message
     // "sparse tensor allocation should not escape function".
     // As such, we insert a trivial tensor convert which will be removed by
LogicalResult matchAndRewrite(OutOp op,
     Value indices = dimSizes; // Reuse the dimSizes buffer for indices.
     Type eltTp = srcTp.getElementType();
     SmallString<18> outNextFuncName{"outSparseTensorWriterNext",
                                     primaryTypeFunctionSuffix(eltTp)};
     Value value = genAllocaScalar(rewriter, loc, eltTp);
     ModuleOp module = op->getParentOfType<ModuleOp>();
     // For each element in the source tensor, output the element.
     rewriter.create<ForeachOp>(
-        loc, src, [&](OpBuilder &builder, Location loc, ValueRange args) {
+        loc, src, llvm::None,
+        [&](OpBuilder &builder, Location loc, ValueRange args, Value v,
+            ValueRange reduc) {
           for (uint64_t i = 0; i < rank; i++) {
             rewriter.create<memref::StoreOp>(loc, args[i], indices,
                                              constantIndex(builder, loc, i));
           }
-          rewriter.create<memref::StoreOp>(loc, args.back(), value);
+          rewriter.create<memref::StoreOp>(loc, v, value);
           SmallVector<Value, 4> operands{writer, rankValue, indices, value};
           FlatSymbolRefAttr fn = getFunc(module, outNextFuncName, {}, operands,
                                          EmitCInterface::On);
           builder.create<func::CallOp>(loc, TypeRange(), fn, operands);
           builder.create<sparse_tensor::YieldOp>(loc);
         });
 
     // Release the writer.

mlir/test/Dialect/SparseTensor/invalid.mlir

func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>) -> () {
   sparse_tensor.foreach in %arg0 : tensor<2x4xf64, #DCSR> do {
     ^bb0(%1: index, %2: index, %v: f32) :
   }
   return
 }
 
 // -----
 
+#DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
+func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>) -> () {
+  // expected-error@+1 {{Unmatched element type between input tensor and block argument}}
+  sparse_tensor.foreach in %arg0 : tensor<2x4xf64, #DCSR> do {
+    ^bb0(%1: index, %2: index, %v: f32) :
+  }
+  return
+}
+
+// -----
+
+#DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
+func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>, %arg1: f32) -> () {
+  // expected-error@+1 {{Mismatch in number of init arguments and results}}
+  sparse_tensor.foreach in %arg0 init(%arg1) : tensor<2x4xf64, #DCSR>, f32 do {
+    ^bb0(%1: index, %2: index, %v: f32, %r1 : i32) :
+  }
+  return
+}
+
+// -----
+
+#DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
+func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>, %arg1: f32) -> () {
+  // expected-error@+1 {{Mismatch in types of init arguments and results}}
+  %1 = sparse_tensor.foreach in %arg0 init(%arg1) : tensor<2x4xf64, #DCSR>, f32 -> i32 do {
+    ^bb0(%1: index, %2: index, %v: f32, %r0 : f32) :
+  }
+  return
+}
+
+// -----
+
+#DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
+func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>, %arg1: f32) -> () {
+  // expected-error@+1 {{Mismatch in types of yield values and results}}
+  %1 = sparse_tensor.foreach in %arg0 init(%arg1) : tensor<2x4xf64, #DCSR>, f32 -> f32 do {
+    ^bb0(%1: index, %2: index, %v: f32, %r0 : f32) :
+      sparse_tensor.yield %1 : index
+  }
+  return
+}
+
+// -----
+
 // TODO: a test case with empty xs doesn't work due to some parser issues.
 
 func.func @sparse_sort_x_type( %arg0: index, %arg1: memref<?xf32>) {
   // expected-error@+1 {{operand #1 must be 1D memref of integer or index values}}
   sparse_tensor.sort %arg0, %arg1: memref<?xf32>
 }
 
 // -----

mlir/test/Dialect/SparseTensor/roundtrip.mlir

 //       CHECK:  ^bb0(%arg1: index, %arg2: index, %arg3: f64):
 func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>) -> () {
   sparse_tensor.foreach in %arg0 : tensor<2x4xf64, #DCSR> do {
     ^bb0(%1: index, %2: index, %v: f64) :
   }
   return
 }
 
+// -----
+
+#DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
+
+// CHECK-LABEL: func @sparse_tensor_foreach(
+//  CHECK-SAME:   %[[A0:.*]]: tensor<2x4xf64, #sparse_tensor.encoding<{{{.*}}}>>, 
+//  CHECK-SAME:   %[[A1:.*]]: f32
+//  CHECK-NEXT:   %[[RET:.*]] = sparse_tensor.foreach in %[[A0]] init(%[[A1]])
+//  CHECK-NEXT:    ^bb0(%[[TMP_1:.*]]: index, %[[TMP_2:.*]]: index, %[[TMP_v:.*]]: f64, %[[TMP_r:.*]]: f32)
+//       CHECK:      sparse_tensor.yield %[[TMP_r]] : f32
+//       CHECK:  }
+func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>, %arg1: f32) -> () {
+  %ret = sparse_tensor.foreach in %arg0 init(%arg1): tensor<2x4xf64, #DCSR>, f32 -> f32
+  do {
+    ^bb0(%1: index, %2: index, %v: f64, %r: f32) : 
+      sparse_tensor.yield %r : f32
+  }
+  return
+}
+
 // ----
 
 // CHECK-LABEL: func @sparse_sort_1d0v(
 //  CHECK-SAME: %[[A:.*]]: index,
 //  CHECK-SAME: %[[B:.*]]: memref<?xindex>)
 //       CHECK: sparse_tensor.sort %[[A]], %[[B]] : memref<?xindex>
 //       CHECK: return %[[B]]
 func.func @sparse_sort_1d0v(%arg0: index, %arg1: memref<?xindex>) -> (memref<?xindex>) {