[mlir][sparse] Introducing a new sparse_tensor.foreach operator.

[aartbik]

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

191

let's make a new "section" for this. e.g. iteration, or something like that

[Peiming]

move into a new section

[Peiming]

rename section

[aartbik]

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

733

in a sparse tensor -> in the given sparse tensor

739

coordination -> coordinates

741

I think you can remove 741/743 since that really only applies to the memory related ops that do something on the side.
It is pretty obvious that a foreach does something more than produce one value ;-)

752

typo: tenosr -> tensor

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

459

static_cast

462

unsigned i

[Peiming]

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

462

surprisingly, getRank() return int64_t

[Peiming]

address comments

[aartbik]

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

462

Oh, but then use int64_t please, also use the idiom

for (int64_t i = 0, e = t.getRank(); i < e; i++)

[Peiming]

Change type in loops

[Peiming]

address comments

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

%0 = sparse_tensor.concatenate %1, %2 { dimension = 0 : index }
        : tensor<64x64xf64, #CSR>, tensor<64x64xf64, #CSR> to tensor<128x64xf64, #CSR>
      ```
    }];
 
   let assemblyFormat = "$inputs attr-dict `:` type($inputs) `to` type($result)";
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse Tensor Management Operations. These operations are "impure" in the
 // sense that they do not properly operate on SSA values. Instead, the behavior
 // is solely defined by side-effects. These operations provide a bridge between
 // "sparsification" on one hand and a support library or actual code generation
 // on the other hand. The semantics of these operations may be refined over time
 // as our sparse abstractions evolve.
 //===----------------------------------------------------------------------===//
 
string description = [{
     ```mlir
     sparse_tensor.out %t, %dest : tensor<1024x1024xf64, #CSR>, !Dest
     ```
   }];
   let assemblyFormat = "$tensor `,` $dest attr-dict `:` type($tensor) `,` type($dest)";
 }
 
 //===----------------------------------------------------------------------===//
-// Sparse Tensor Custom Linalg.Generic Operations.
+// Sparse Tensor Syntax Operations.
 //===----------------------------------------------------------------------===//
 
 def SparseTensor_BinaryOp : SparseTensor_Op<"binary", [NoSideEffect]>,
     Arguments<(ins AnyType:$x, AnyType:$y, UnitAttr:$left_identity, UnitAttr:$right_identity)>,
     Results<(outs AnyType:$output)> {
   let summary = "Binary set operation utilized within linalg.generic";
   let description = [{
       Defines a computation within a `linalg.generic` operation that takes two
def SparseTensor_SelectOp : SparseTensor_Op<"select", [NoSideEffect, SameOperandsAndResultType]>,
   let regions = (region SizedRegion<1>:$region);
   let assemblyFormat = [{
          $x attr-dict `:` type($x) $region
   }];
   let hasVerifier = 1;
 }
 
 def SparseTensor_YieldOp : SparseTensor_Op<"yield", [NoSideEffect, Terminator]>,
-    Arguments<(ins AnyType:$result)> {
+    Arguments<(ins Optional<AnyType>:$result)> {
   let summary = "Yield from sparse_tensor set-like operations";
   let description = [{
-      Yields a value from within a `binary`, `unary`, `reduce`,
-      or `select` block.
+      Yields a value from within a `binary`, `unary`, `reduce`, 
+      `select` or `foreach` block.
 
       Example:
 
       ```
       %0 = sparse_tensor.unary %a : i64 to i64 {
         ^bb0(%arg0: i64):
           %cst = arith.constant 1 : i64
           %ret = arith.addi %arg0, %cst : i64
           sparse_tensor.yield %ret : i64
       }
       ```
   }];
 
+  let builders = [
+    OpBuilder<(ins),
+    [{
+      build($_builder, $_state, Value());
+    }]>
+  ];
+
   let assemblyFormat = [{
         $result attr-dict `:` type($result)
   }];
   let hasVerifier = 1;
 }
 
+def SparseTensor_ForeachOp : SparseTensor_Op<"foreach",
+    [SingleBlockImplicitTerminator<"YieldOp">]>,
+    Arguments<(ins AnySparseTensor:$tensor)>{
+  let summary = "Iterates over non-zero elements in a sparse tensor";
+  let description = [{
+     Iterates over every non-zero element in the given sparse tensor and executes
+     the block.
+
+     For a input sparse 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
+     sparse tensor's element type, representing the actual value loaded from the input
+     tensor at the given coordinates.
+
+     Example:
+     
+     ```mlir
+     sparse_tensor.foreach in %0 : tensor<?x?xf64, #DCSR> do {
+      ^bb0(%arg1: index, %arg2: index, %arg3: f64):
+        do something...
+     }
+     ```
+  }];
+  
+  let regions = (region AnyRegion:$region);
+  let assemblyFormat = "`in` $tensor attr-dict `:` type($tensor)  `do` $region";
+  let hasVerifier = 1;
+}
+
 #endif // SPARSETENSOR_OPS

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

if (typ != inputTypes[i])
       return op->emitError() << regionName << " region argument " << (i + 1)
                              << " type mismatch";
   }
   Operation *term = region.front().getTerminator();
   YieldOp yield = dyn_cast<YieldOp>(term);
   if (!yield)
     return op->emitError() << regionName
                            << " region must end with sparse_tensor.yield";
-  if (yield.getOperand().getType() != outputType)
+  if (!yield.getResult() || yield.getResult().getType() != outputType)
     return op->emitError() << regionName << " region yield type mismatch";
 
   return success();
 }
 
 LogicalResult BinaryOp::verify() {
   NamedAttrList attrs = (*this)->getAttrs();
   Type leftType = getX().getType();
if (i == concatDim) {
         prev = d;
       }
     }
   }
 
   return success();
 }
 
+LogicalResult ForeachOp::verify() {
+  auto t = getTensor().getType().cast<RankedTensorType>();
+  auto args = getBody()->getArguments();
+
+  if (static_cast<unsigned>(t.getRank()) + 1 != args.size())
+    return emitError("Unmatched number of arguments in the block");
+
+  for (int i = 0; i < t.getRank(); 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();
+  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() {
   Type inputType = getX().getType();
   LogicalResult regionResult = success();
 
   // Check correct number of block arguments and return type.
   Region &formula = getRegion();
   regionResult = verifyNumBlockArgs(this, formula, "reduce",
                                     TypeRange{inputType, inputType}, inputType);
LogicalResult SelectOp::verify() {
 
   return success();
 }
 
 LogicalResult YieldOp::verify() {
   // Check for compatible parent.
   auto *parentOp = (*this)->getParentOp();
   if (isa<BinaryOp>(parentOp) || isa<UnaryOp>(parentOp) ||
-      isa<ReduceOp>(parentOp) || isa<SelectOp>(parentOp))
+      isa<ReduceOp>(parentOp) || isa<SelectOp>(parentOp) ||
+      isa<ForeachOp>(parentOp))
     return success();
 
   return emitOpError("expected parent op to be sparse_tensor unary, binary, "
-                     "reduce, or select");
+                     "reduce, select or foreach");
 }
 
 //===----------------------------------------------------------------------===//
 // TensorDialect Methods.
 //===----------------------------------------------------------------------===//
 
 void SparseTensorDialect::initialize() {
   addAttributes<

mlir/test/Dialect/SparseTensor/invalid.mlir

func.func @invalid_concat_size_mismatch(%arg0: tensor<2x4xf64, #DC>,
                                         %arg2: tensor<4x4xf64, #DC>) -> tensor<9x4xf64, #DC> {
   // expected-error@+1 {{All dimensions (expect for the concatenating one) should be equal}}
   %0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index}
        : tensor<2x4xf64, #DC>,
          tensor<3x3xf64, #DC>,
          tensor<4x4xf64, #DC> to tensor<9x4xf64, #DC>
   return %0 : tensor<9x4xf64, #DC>
 }
+
+// -----
+
+#DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
+func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>) -> () {
+  // expected-error@+1 {{Unmatched number of arguments in the block}}
+  sparse_tensor.foreach in %arg0 : tensor<2x4xf64, #DCSR> do {
+    ^bb0(%1: index, %2: index, %3: index, %v: f64) :
+  }
+  return
+}
+
+// -----
+
+#DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
+func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>) -> () {
+  // expected-error@+1 {{Expecting Index type for argument at index 1}}
+  sparse_tensor.foreach in %arg0 : tensor<2x4xf64, #DCSR> do {
+    ^bb0(%1: index, %2: f64, %v: f64) :
+  }
+  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
+}

mlir/test/Dialect/SparseTensor/roundtrip.mlir

func.func @concat_sparse_sparse(%arg0: tensor<2x4xf64, #SparseMatrix>,
                                 %arg1: tensor<3x4xf64, #SparseMatrix>,
                                 %arg2: tensor<4x4xf64, #SparseMatrix>) -> tensor<9x4xf64, #SparseMatrix> {
   %0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index}
        : tensor<2x4xf64, #SparseMatrix>,
          tensor<3x4xf64, #SparseMatrix>,
          tensor<4x4xf64, #SparseMatrix> to tensor<9x4xf64, #SparseMatrix>
   return %0 : tensor<9x4xf64, #SparseMatrix>
 }
+
+// -----
+
+#DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
+
+// CHECK-LABEL: func @sparse_tensor_foreach(
+//  CHECK-SAME: %[[A0:.*]]: tensor<2x4xf64
+//       CHECK: sparse_tensor.foreach in %[[A0]] : 
+//       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
+}