[MLIR][Shape] Allow `shape.reduce` to operate on extent tensors

[frgossen]

Summary

Allow shape.reduce to take both shape.shape and tensor<?xindex> as an
argument.

[herhut]

Nice, thank you!

mlir/test/Dialect/Shape/invalid.mlir

25–26

Does this still pass? You have changed the error messages. Also, please add another version where the expected type is IndexType.

[frgossen]

Correct test case

mlir/test/Dialect/Shape/invalid.mlir

25–26

It passes because it's a substring and that is enough.
You're right though, it should be changed.

[herhut]

Thanks!

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

def Shape_NumElementsOp : Shape_Op<"num_elements", [NoSideEffect]> {
 
   let assemblyFormat = "$shape attr-dict";
 
   let hasFolder = 1;
 }
 
 def Shape_ReduceOp : Shape_Op<"reduce",
     [SingleBlockImplicitTerminator<"YieldOp">]> {
-  let summary = "Returns an expression reduced over a shape";
+  let summary = "Returns an expression reduced over a shape or extent tensor";
   let description = [{
-    An operation that takes as input a shape, number of initial values and has a
-    region/function that is applied repeatedly for every dimension of the shape.
+    An operation that takes as input a shape or extent tensor, and a number of
+    initial values. This operation has a region/function that is applied
+    repeatedly for every extent of the input. Starting with the initial values,
+    the individual extents are then aggregated as defined by the associated
+    region.
 
     Conceptually this op performs the following reduction:
 
     ```
     res[] = init;
-    for (int i = 0, e = shape.rank(); i != e; ++i) {
+    for (int i = 0, i < shape.rank(); i++) {
       res = fn(i, shape[i], res[0], ..., res[n]);
     }
     ```
 
-    Where fn is provided by the user and the result of the reduce op is the
+    Where `fn` is provided by the user and the result of the reduce op is the
     last computed output of the reduce function. As an example, computing the
-    number of elements
+    number of elements can be defined as follows:
 
     ```mlir
     func @reduce(%shape : !shape.shape, %init : !shape.size) -> !shape.size {
       %num_elements = shape.reduce(%shape, %init) -> !shape.size  {
         ^bb0(%index: index, %dim: !shape.size, %acc: !shape.size):
           %updated_acc = "shape.mul"(%acc, %dim) :
             (!shape.size, !shape.size) -> !shape.size
           shape.yield %updated_acc : !shape.size
       }
       return %num_elements : !shape.size
     }
     ```
-
-    If the shape is unranked, then the results of the op is also unranked.
   }];
 
-  let arguments = (ins Shape_ShapeType:$shape, Variadic<AnyType>:$initVals);
+  let arguments = (ins Shape_ShapeOrExtentTensorType:$shape,
+                       Variadic<AnyType>:$initVals);
   let results = (outs Variadic<AnyType>:$result);
   let regions = (region SizedRegion<1>:$region);
 
   let builders = [
     OpBuilder<"OpBuilder &builder, OperationState &result, "
               "Value shape, ValueRange initVals">,
   ];
 

mlir/lib/Dialect/Shape/IR/Shape.cpp

for (Type initValType : initVals.getTypes()) {
     result.addTypes(initValType);
   }
 }
 
 static LogicalResult verify(ReduceOp op) {
   // Verify block arg types.
   Block &block = op.region().front();
 
+  // The block takes index, extent, and aggregated values as arguments.
   auto blockArgsCount = op.initVals().size() + 2;
   if (block.getNumArguments() != blockArgsCount)
     return op.emitOpError() << "ReduceOp body is expected to have "
                             << blockArgsCount << " arguments";
 
-  if (block.getArgument(0).getType() != IndexType::get(op.getContext()))
+  // The first block argument is the index and must always be of type `index`.
+  if (!block.getArgument(0).getType().isa<IndexType>())
     return op.emitOpError(
         "argument 0 of ReduceOp body is expected to be of IndexType");
 
-  if (block.getArgument(1).getType() != SizeType::get(op.getContext()))
+  // The second block argument is the extent and must be of type `size` or
+  // `index`, depending on whether the reduce operation is applied to a shape or
+  // to an extent tensor.
+  Type extentTy = block.getArgument(1).getType();
+  if (op.shape().getType().isa<ShapeType>()) {
+    if (!extentTy.isa<SizeType>())
+      return op.emitOpError("argument 1 of ReduceOp body is expected to be of "
+                            "SizeType if the ReduceOp operates on a ShapeType");
+  } else {
+    if (!extentTy.isa<IndexType>())
       return op.emitOpError(
-        "argument 1 of ReduceOp body is expected to be of SizeType");
+          "argument 1 of ReduceOp body is expected to be of IndexType if the "
+          "ReduceOp operates on an extent tensor");
+  }
 
   for (auto type : llvm::enumerate(op.initVals()))
     if (block.getArgument(type.index() + 2).getType() != type.value().getType())
       return op.emitOpError()
              << "type mismatch between argument " << type.index() + 2
              << " of ReduceOp body and initial value " << type.index();
   return success();
 }
 
 static ParseResult parseReduceOp(OpAsmParser &parser, OperationState &result) {
-  auto *ctx = parser.getBuilder().getContext();
   // Parse operands.
   SmallVector<OpAsmParser::OperandType, 3> operands;
+  Type shapeOrExtentTensorType;
   if (parser.parseOperandList(operands, /*requiredOperandCount=*/-1,
                               OpAsmParser::Delimiter::Paren) ||
+      parser.parseColonType(shapeOrExtentTensorType) ||
       parser.parseOptionalArrowTypeList(result.types))
     return failure();
 
   // Resolve operands.
   auto initVals = llvm::makeArrayRef(operands).drop_front();
-  if (parser.resolveOperand(operands.front(), ShapeType::get(ctx),
+  if (parser.resolveOperand(operands.front(), shapeOrExtentTensorType,
                             result.operands) ||
       parser.resolveOperands(initVals, result.types, parser.getNameLoc(),
                              result.operands))
     return failure();
 
   // Parse the body.
   Region *body = result.addRegion();
   if (parser.parseRegion(*body, /*args=*/{}, /*argTypes=*/{}))
     return failure();
 
   // Parse attributes.
   if (parser.parseOptionalAttrDict(result.attributes))
     return failure();
 
   return success();
 }
 
 static void print(OpAsmPrinter &p, ReduceOp op) {
   p << op.getOperationName() << '(' << op.shape() << ", " << op.initVals()
-    << ") ";
+    << ") : " << op.shape().getType();
   p.printOptionalArrowTypeList(op.getResultTypes());
   p.printRegion(op.region());
   p.printOptionalAttrDict(op.getAttrs());
 }
 
 namespace mlir {
 namespace shape {
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Shape/IR/ShapeOps.cpp.inc"
 
 } // namespace shape
 } // namespace mlir

mlir/test/Conversion/ShapeToSCF/shape-to-scf.mlir

 // RUN: mlir-opt -convert-shape-to-scf -split-input-file %s | FileCheck %s
 
-// CHECK-LABEL: shape_reduce
-// CHECK-SAME:   [[SHAPE:%.*]]: !shape.shape) -> !shape.size {
+// CHECK-LABEL: @shape_reduce
+// CHECK-SAME:  ([[SHAPE:%.*]]: !shape.shape) -> !shape.size
 func @shape_reduce(%shape : !shape.shape) -> !shape.size {
   %init = shape.const_size 1
-  %num_elements = shape.reduce(%shape, %init) -> !shape.size {
+  %num_elements = shape.reduce(%shape, %init) : !shape.shape -> !shape.size {
     ^bb0(%index: index, %dim: !shape.size, %acc: !shape.size):
       %new_acc = shape.mul %acc, %dim
       shape.yield %new_acc : !shape.size
   }
   return %num_elements : !shape.size
 }
 // CHECK-NEXT: [[SHAPE_C1:%.*]] = shape.const_size 1
 // CHECK-NEXT: [[C0:%.*]] = constant 0 : index

mlir/test/Dialect/Shape/invalid.mlir

 // RUN: mlir-opt %s -split-input-file -verify-diagnostics
 
 func @reduce_op_args_num_mismatch(%shape : !shape.shape, %init : !shape.size) {
   // expected-error@+1 {{ReduceOp body is expected to have 3 arguments}}
-  %num_elements = shape.reduce(%shape, %init) -> !shape.size {
+  %num_elements = shape.reduce(%shape, %init) : !shape.shape -> !shape.size {
     ^bb0(%index: index, %dim: !shape.size):
       shape.yield %dim : !shape.size
   }
 }
 
 // -----
 
 func @reduce_op_arg0_wrong_type(%shape : !shape.shape, %init : !shape.size) {
   // expected-error@+1 {{argument 0 of ReduceOp body is expected to be of IndexType}}
-  %num_elements = shape.reduce(%shape, %init) -> !shape.size {
+  %num_elements = shape.reduce(%shape, %init) : !shape.shape -> !shape.size {
     ^bb0(%index: f32, %dim: !shape.size, %acc: !shape.size):
       %new_acc = "shape.add"(%acc, %dim)
           : (!shape.size, !shape.size) -> !shape.size
       shape.yield %new_acc : !shape.size
   }
 }
 
 // -----
 
 func @reduce_op_arg1_wrong_type(%shape : !shape.shape, %init : !shape.size) {
-  // expected-error@+1 {{argument 1 of ReduceOp body is expected to be of SizeType}}
-  %num_elements = shape.reduce(%shape, %init) -> !shape.size {
+  // expected-error@+1 {{argument 1 of ReduceOp body is expected to be of SizeType if the ReduceOp operates on a ShapeType}}
+  %num_elements = shape.reduce(%shape, %init) : !shape.shape -> !shape.size {
     ^bb0(%index: index, %dim: f32, %lci: !shape.size):
       shape.yield
   }
 }
 
 // -----
 
+func @reduce_op_arg1_wrong_type(%shape : tensor<?xindex>, %init : index) {
+  // expected-error@+1 {{argument 1 of ReduceOp body is expected to be of IndexType if the ReduceOp operates on an extent tensor}}
+  %num_elements = shape.reduce(%shape, %init) : tensor<?xindex> -> index {
+    ^bb0(%index: index, %dim: f32, %lci: index):
+      shape.yield
+  }
+}
+
+// -----
+
 func @reduce_op_init_type_mismatch(%shape : !shape.shape, %init : f32) {
   // expected-error@+1 {{type mismatch between argument 2 of ReduceOp body and initial value 0}}
-  %num_elements = shape.reduce(%shape, %init) -> f32 {
+  %num_elements = shape.reduce(%shape, %init) : !shape.shape -> f32 {
     ^bb0(%index: index, %dim: !shape.size, %lci: !shape.size):
       shape.yield
   }
 }
 
 // -----
 
 func @yield_op_args_num_mismatch(%shape : !shape.shape, %init : !shape.size) {
   // expected-error@+3 {{number of operands does not match number of results of its parent}}
-  %num_elements = shape.reduce(%shape, %init) -> !shape.size {
+  %num_elements = shape.reduce(%shape, %init) : !shape.shape -> !shape.size {
     ^bb0(%index: index, %dim: !shape.size, %lci: !shape.size):
       shape.yield %dim, %dim : !shape.size, !shape.size
   }
 }
 
 // -----
 
 func @yield_op_type_mismatch(%shape : !shape.shape, %init : !shape.size) {
   // expected-error@+4 {{types mismatch between yield op and its parent}}
-  %num_elements = shape.reduce(%shape, %init) -> !shape.size {
+  %num_elements = shape.reduce(%shape, %init) : !shape.shape -> !shape.size {
     ^bb0(%index: index, %dim: !shape.size, %lci: !shape.size):
       %c0 = constant 1 : index
       shape.yield %c0 : index
   }
 }
 
 // -----
 
 func @assuming_all_op_too_few_operands() {
   // expected-error@+1 {{no operands specified}}
   %w0 = shape.assuming_all
   return
 }

mlir/test/Dialect/Shape/ops.mlir

 // RUN: mlir-opt -split-input-file %s | mlir-opt | FileCheck %s
 // Verify the printed output can be parsed.
 // RUN: mlir-opt %s | mlir-opt | FileCheck %s
 // Verify the generic form can be parsed.
 // RUN: mlir-opt -mlir-print-op-generic %s | mlir-opt | FileCheck %s
 
 // CHECK-LABEL: shape_num_elements
 func @shape_num_elements(%shape : !shape.shape) -> !shape.size {
-  %init = shape.const_size 0
-  %num_elements = shape.reduce(%shape, %init) -> !shape.size {
-    ^bb0(%index: index, %dim: !shape.size, %lci: !shape.size):
-      %acc = shape.add %lci, %dim
-      shape.yield %acc : !shape.size
+  %init = shape.const_size 1
+  %num_elements = shape.reduce(%shape, %init) : !shape.shape -> !shape.size {
+    ^bb0(%index : index, %extent : !shape.size, %acc : !shape.size):
+      %acc_next = shape.mul %acc, %extent
+      shape.yield %acc_next : !shape.size
   }
   return %num_elements : !shape.size
 }
 
+// CHECK-LABEL: extent_tensor_num_elements
+func @extent_tensor_num_elements(%shape : tensor<?xindex>) -> index {
+  %init = constant 1 : index
+  %num_elements = shape.reduce(%shape, %init) : tensor<?xindex> -> index {
+    ^bb0(%index : index, %extent : index, %acc : index):
+      %acc_next = muli %acc, %extent : index
+      shape.yield %acc_next : index
+  }
+  return %num_elements : index
+}
+
 func @test_shape_num_elements_unknown() {
   %0 = "shape.unknown_shape"() : () -> !shape.shape
   %1 = call @shape_num_elements(%0) : (!shape.shape) -> (!shape.size)
   %2 = "shape.print"(%1) : (!shape.size) -> !shape.size
   return
 }
 
 func @test_shape_num_elements_fixed() {

mlir/test/Dialect/Shape/shape-to-shape.mlir

 // RUN: mlir-opt -shape-to-shape-lowering -split-input-file %s | FileCheck %s
 
 // CHECK-LABEL: func @num_elements_to_reduce(
-// CHECK-SAME:    [[ARG:%.*]]: !shape.shape) -> [[SIZE_TY:!.*]] {
+// CHECK-SAME:    [[ARG:%.*]]: !shape.shape) -> !shape.size {
 func @num_elements_to_reduce(%shape : !shape.shape) -> !shape.size {
   %num_elements = shape.num_elements %shape
   return %num_elements : !shape.size
 }
 // CHECK: [[C1:%.*]] = shape.const_size 1
-// CHECK: [[NUM_ELEMENTS:%.*]] = shape.reduce([[ARG]], [[C1]])  -> [[SIZE_TY]]
-// CHECK: ^bb0({{.*}}: index, [[DIM:%.*]]: [[SIZE_TY]], [[ACC:%.*]]: [[SIZE_TY]]
+// CHECK: [[NUM_ELEMENTS:%.*]] = shape.reduce([[ARG]], [[C1]]) : !shape.shape -> !shape.size
+// CHECK: ^bb0({{.*}}: index, [[DIM:%.*]]: !shape.size, [[ACC:%.*]]: !shape.size
 // CHECK:   [[NEW_ACC:%.*]] = shape.mul [[DIM]], [[ACC]]
-// CHECK:   shape.yield [[NEW_ACC]] : [[SIZE_TY]]
+// CHECK:   shape.yield [[NEW_ACC]] : !shape.size
 // CHECK: }
-// CHECK: return [[NUM_ELEMENTS]] : [[SIZE_TY]]
+// CHECK: return [[NUM_ELEMENTS]] : !shape.size