[mlir][sparse] Add a new reduce operation in sparse_tensor dialect

[jim22k]

Summary

Custom reduction is a binary operators, but differs from the existing
sparse_tensor.binary operation in significant ways.

1. It does not deal with sparsity overlaps. It is given a finite set of non-empty values and must collapse them to a single value
2. The nature of reduction requires a starting value (or identity value). This is normally zero for summation, but for the product needs to be 1.0. The custom code needs to explicitly pass an identity value.

[aartbik]

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

554

The documentation can use a bit more explanation on why we added this.

For example, you can show that "regular" reductions can be expressed "the easy way", but this makes the sparse specifications more expressive (I want to avoid that newcomers think *all* reductions need to be defined this way)

560

empty line after Example:

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

378–379

since we don't use the result of the cast, why not simply use a "isa"?

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp

959 ↗

(On Diff #437689)

can't we start a new == 3 section then for ternary operations?

mlir/test/Dialect/SparseTensor/invalid.mlir

344

looks like we never had an invalid test for "must end with sparse_tensor.yield", for this or another one, so perhaps add this as part of this revision to get full coverage

[jim22k]

• Merge with latest changes (D128000)
• Update based on feedback

[jim22k]

Fix formatting

[aartbik]

Perhaps for the sake of progress, can we split this into two revisions? Once that introduces the new sparse dialect op, and then one that provides the right implementation?
I feel that the former is almost ready to go in, but the latter needs some work, so splitting this up into smaller pieces may make it easier.

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

570

This example is missing the surrounding linalg.generic, making it also harder to see what %x and %y are exactly. Also, in the examples above you use %a,%b,%c for the top level, and then %arg0 again in the block inside the semi-ring op. Please make the syntax in all examples similar, so it is easier to read and recognize common components.

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

745 ↗

(On Diff #440630)

I am puzzled we need to have new code in the merger for this, can't we just handle that in the sparse codegen file as before. Also, please note that we need to deal with vector values during vectorization. Probably not for the semi ring op, but your ReducID seems to deal with the other cases too.

Since you are using so many if (codegen.redKind == kNoReduc) tests, it feels like you are entering the wrong methods, and should have deal with reductions at a higher level (see updateReduc() and friends)

[jim22k]

Add reduce operation definition and validation tests

[jim22k]

I split the PR in two. This is now just the .td definition and validation tests. Hopefully this makes it easier to approve, and we can discuss implementation in a subsequent PR.

[aartbik]

In D128004#3652267, @jim22k wrote:

I split the PR in two. This is now just the .td definition and validation tests. Hopefully this makes it easier to approve, and we can discuss implementation in a subsequent PR.

Thanks for splitting this up! This part is directly good to go, and we can start working on the implementation with a smaller revision!

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

def SparseTensor_UnaryOp : SparseTensor_Op<"unary", [NoSideEffect]>,
   let assemblyFormat = [{
         $x attr-dict `:` type($x) `to` type($output) `\n`
         `present` `=` $presentRegion `\n`
         `absent` `=` $absentRegion
   }];
   let hasVerifier = 1;
 }
 
+def SparseTensor_ReduceOp : SparseTensor_Op<"reduce", [NoSideEffect, SameOperandsAndResultType]>,
+    Arguments<(ins AnyType:$x, AnyType:$y, AnyType:$identity)>,
+    Results<(outs AnyType:$output)> {
+  let summary = "Custom reduction operation utilized within linalg.generic";
+  let description = [{
+      Defines a computation with a `linalg.generic` operation that takes two
+      operands and an identity value and reduces all values down to a single
+      result based on the computation in the region.
+
+      The region must contain exactly one block taking two arguments. The block
+      must end with a sparse_tensor.yield and the output must match the input
+      argument types.
+
+      Note that this operation is only required for custom reductions beyond the
+
+      standard operations (add, mul, and, or, etc). The `linalg.generic`
+      `iterator_types` defines which indices are being reduced. When the associated
+      operands are used in an operation, a reduction will occur. The use of this
+      explicit `reduce` operation is not required in most cases.
+
+      Example of Matrix->Vector reduction using max(product(x_i), 100):
+
+      ```mlir
+      %cf1 = arith.constant 1.0 : f64
+      %cf100 = arith.constant 100.0 : f64
+      %C = bufferization.alloc_tensor...
+      %0 = linalg.generic #trait
+        ins(%A: tensor<?x?xf64, #SparseMatrix>)
+        outs(%C: tensor<?xf64, #SparseVec>) {
+        ^bb0(%a: f64, %c: f64) :
+          %result = sparse_tensor.reduce %c, %a, %cf1 : f64 {
+              ^bb0(%arg0: f64, %arg1: f64):
+                %0 = arith.mulf %arg0, %arg1 : f64
+                %cmp = arith.cmpf "ogt", %0, %cf100 : f64
+                %ret = arith.select %cmp, %cf100, %0 : f64
+                sparse_tensor.yield %ret : f64
+            }
+          linalg.yield %result : f64
+      } -> tensor<?xf64, #SparseVec>
+      ```
+  }];
+
+  let regions = (region SizedRegion<1>:$region);
+
+  let assemblyFormat = [{
+         $x `,` $y `,` $identity attr-dict `:` type($output) $region
+  }];
+  let hasVerifier = 1;
+}
+
 def SparseTensor_YieldOp : SparseTensor_Op<"yield", [NoSideEffect, Terminator]>,
     Arguments<(ins AnyType:$result)> {
   let summary = "Yield from sparse_tensor set-like operations";
   let description = [{
       Yields a value from within a `binary` or `unary` block.
 
       Example:
 

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

regionResult =
         verifyNumBlockArgs(this, absent, "absent", TypeRange{}, outputType);
     if (failed(regionResult))
       return regionResult;
   }
 
   return success();
 }
 
+LogicalResult ReduceOp::verify() {
+  Type inputType = x().getType();
+  LogicalResult regionResult = success();
+
+  // Check correct number of block arguments and return type.
+  Region &formula = region();
+  if (!formula.empty()) {
+    regionResult = verifyNumBlockArgs(
+        this, formula, "reduce", TypeRange{inputType, inputType}, inputType);
+    if (failed(regionResult))
+      return regionResult;
+  }
+
+  return success();
+}
+
 LogicalResult YieldOp::verify() {
   // Check for compatible parent.
   auto *parentOp = (*this)->getParentOp();
-  if (auto binaryOp = dyn_cast<BinaryOp>(parentOp))
-    return success();
-  if (auto unaryOp = dyn_cast<UnaryOp>(parentOp))
+  if (isa<BinaryOp>(parentOp) || isa<UnaryOp>(parentOp) ||
+      isa<ReduceOp>(parentOp))
     return success();
 
-  return emitOpError("expected parent op to be sparse_tensor binary or unary");
+  return emitOpError(
+      "expected parent op to be sparse_tensor unary, binary, or reduce");
 }
 
 //===----------------------------------------------------------------------===//
 // TensorDialect Methods.
 //===----------------------------------------------------------------------===//
 
 void SparseTensorDialect::initialize() {
   addAttributes<

mlir/test/Dialect/SparseTensor/invalid.mlir

%0 = sparse_tensor.binary %arg0, %arg1 : i64, f64 to f64
     overlap={}
     left=identity
     right=identity
   return %0 : f64
 }
 
 // -----
 
+func.func @invalid_binary_wrong_yield(%arg0: f64, %arg1: f64) -> f64 {
+  // expected-error@+1 {{left region must end with sparse_tensor.yield}}
+  %0 = sparse_tensor.binary %arg0, %arg1 : f64, f64 to f64
+    overlap={}
+    left={
+      ^bb0(%x: f64):
+        tensor.yield %x : f64
+    }
+    right=identity
+  return %0 : f64
+}
+
+// -----
+
 func.func @invalid_unary_argtype_mismatch(%arg0: f64) -> f64 {
   // expected-error@+1 {{present region argument 1 type mismatch}}
   %r = sparse_tensor.unary %arg0 : f64 to f64
     present={
       ^bb0(%x: index):
         sparse_tensor.yield %x : index
     }
     absent={}
%0 = sparse_tensor.unary %arg0 : f64 to f64
     present={
       ^bb0(%x: f64):
         %1 = arith.constant 0.0 : f32
         sparse_tensor.yield %1 : f32
     }
     absent={}
   return %0 : f64
 }
+
+// -----
+
+func.func @invalid_unary_wrong_yield(%arg0: f64) -> f64 {
+  // expected-error@+1 {{present region must end with sparse_tensor.yield}}
+  %0 = sparse_tensor.unary %arg0 : f64 to f64
+    present={
+      ^bb0(%x: f64):
+        tensor.yield %x : f64
+    }
+    absent={}
+  return %0 : f64
+}
+
+// -----
+
+func.func @invalid_reduce_num_args_mismatch(%arg0: f64, %arg1: f64) -> f64 {
+  %cf1 = arith.constant 1.0 : f64
+  // expected-error@+1 {{reduce region must have exactly 2 arguments}}
+  %r = sparse_tensor.reduce %arg0, %arg1, %cf1 : f64 {
+      ^bb0(%x: f64):
+        sparse_tensor.yield %x : f64
+    }
+  return %r : f64
+}
+
+// -----
+
+func.func @invalid_reduce_block_arg_type_mismatch(%arg0: i64, %arg1: i64) -> i64 {
+  %ci1 = arith.constant 1 : i64
+  // expected-error@+1 {{reduce region argument 1 type mismatch}}
+  %r = sparse_tensor.reduce %arg0, %arg1, %ci1 : i64 {
+      ^bb0(%x: f64, %y: f64):
+        %cst = arith.constant 2 : i64
+        sparse_tensor.yield %cst : i64
+    }
+  return %r : i64
+}
+
+// -----
+
+func.func @invalid_reduce_return_type_mismatch(%arg0: f64, %arg1: f64) -> f64 {
+  %cf1 = arith.constant 1.0 : f64
+  // expected-error@+1 {{reduce region yield type mismatch}}
+  %r = sparse_tensor.reduce %arg0, %arg1, %cf1 : f64 {
+      ^bb0(%x: f64, %y: f64):
+        %cst = arith.constant 2 : i64
+        sparse_tensor.yield %cst : i64
+    }
+  return %r : f64
+}
+
+// -----
+
+func.func @invalid_reduce_wrong_yield(%arg0: f64, %arg1: f64) -> f64 {
+  %cf1 = arith.constant 1.0 : f64
+  // expected-error@+1 {{reduce region must end with sparse_tensor.yield}}
+  %r = sparse_tensor.reduce %arg0, %arg1, %cf1 : f64 {
+      ^bb0(%x: f64, %y: f64):
+        %cst = arith.constant 2 : i64
+        tensor.yield %cst : i64
+    }
+  return %r : f64
+}

mlir/test/Dialect/SparseTensor/roundtrip.mlir

%r = sparse_tensor.unary %arg0 : f64 to i64
     present={
       ^bb0(%x: f64):
         %ret = arith.fptosi %x : f64 to i64
         sparse_tensor.yield %ret : i64
     }
     absent={}
   return %r : i64
 }
+
+// -----
+
+#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
+
+// CHECK-LABEL: func @sparse_reduce_2d_to_1d(
+//  CHECK-SAME:   %[[A:.*]]: f64, %[[B:.*]]: f64) -> f64 {
+//       CHECK:   %[[Z:.*]] = arith.constant 0.000000e+00 : f64
+//       CHECK:   %[[C1:.*]] = sparse_tensor.reduce %[[A]], %[[B]], %[[Z]] : f64 {
+//       CHECK:       ^bb0(%[[A1:.*]]: f64, %[[B1:.*]]: f64):
+//       CHECK:         sparse_tensor.yield %[[A1]] : f64
+//       CHECK:     }
+//       CHECK:   return %[[C1]] : f64
+//       CHECK: }
+func.func @sparse_reduce_2d_to_1d(%arg0: f64, %arg1: f64) -> f64 {
+  %cf0 = arith.constant 0.0 : f64
+  %r = sparse_tensor.reduce %arg0, %arg1, %cf0 : f64 {
+      ^bb0(%x: f64, %y: f64):
+        sparse_tensor.yield %x : f64
+    }
+  return %r : f64
+}
 No newline at end of file