[MLIR][Shape] Allow `shape.mul` to operate in indices

[frgossen]

Summary

Depends On D84436

[frgossen]

Rebase

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

let description = [{
     ```
   }];
 
   let arguments = (ins Shape_ShapeOrSizeType:$arg0, Shape_ShapeOrSizeType:$arg1,
                    OptionalAttr<StrAttr>:$error);
   let results = (outs Shape_ShapeOrSizeType:$result);
 }
 
-def Shape_MulOp : Shape_Op<"mul", [Commutative, SameOperandsAndResultType]> {
-  let summary = "Multiplication of sizes";
+def Shape_MulOp : Shape_Op<"mul", [Commutative, NoSideEffect]> {
+  let summary = "Multiplication of sizes and indices";
   let description = [{
-    Multiplies two valid sizes as follows:
-    - lhs * rhs = unknown if either lhs or rhs unknown;
-    - lhs * rhs = (int)lhs * (int)rhs if both known;
+    Multiplies two sizes or indices. If either operand is an error it will be
+    propagated to the result. The operands can be of type `size` or `index`. If
+    at least one of the operands can hold an error, i.e. if it is of type `size`,
+    then also the result must be of type `size`. If error propagation is not
+    possible because both operands are of type `index` then the result must also
+    be of type `index`.
   }];
 
-  let arguments = (ins Shape_SizeType:$lhs, Shape_SizeType:$rhs);
-  let results = (outs Shape_SizeType:$result);
+  let arguments = (ins Shape_SizeOrIndexType:$lhs, Shape_SizeOrIndexType:$rhs);
+  let results = (outs Shape_SizeOrIndexType:$result);
 
-  let assemblyFormat = "$lhs `,` $rhs attr-dict";
+  let assemblyFormat = [{
+    $lhs `,` $rhs `:` type($lhs) `,` type($rhs) `->` type($result) attr-dict
+  }];
+
+  let verifier = [{ return ::verify(*this); }];
 }
 
 def Shape_NumElementsOp : Shape_Op<"num_elements", [NoSideEffect]> {
   let summary = "Returns the number of elements for a given shape";
   let description = [{
     Returns the number of elements for a given shape which is the product of its
     dimensions.
 

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

 namespace {
 #include "ShapeCanonicalization.inc"
 }
 
 static RankedTensorType getExtentTensorType(MLIRContext *ctx) {
   return RankedTensorType::get({ShapedType::kDynamicSize}, IndexType::get(ctx));
 }
 
+static bool isErrorPropagationPossible(ArrayRef<Type> operandTypes) {
+  for (Type ty : operandTypes)
+    if (ty.isa<SizeType>() || ty.isa<ShapeType>() || ty.isa<ValueShapeType>())
+      return true;
+  return false;
+}
+
 ShapeDialect::ShapeDialect(MLIRContext *context)
     : Dialect(getDialectNamespace(), context) {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Shape/IR/ShapeOps.cpp.inc"
       >();
   addTypes<ComponentType, ElementType, ShapeType, SizeType, ValueShapeType,
            WitnessType>();

 //===----------------------------------------------------------------------===//
 // GetExtentOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(GetExtentOp op) {
   Type shapeTy = op.shape().getType();
   Type dimTy = op.dim().getType();
   Type extentTy = op.extent().getType();
-  bool errorPropagationPossible =
-      shapeTy.isa<ShapeType>() || dimTy.isa<SizeType>();
-  if (errorPropagationPossible) {
+  if (isErrorPropagationPossible({shapeTy, dimTy})) {
     if (!extentTy.isa<SizeType>())
       op.emitError()
           << "if at least one of the operands can hold error values then the "
              "result must be of type `size` to propagate them";
   } else {
     if (extentTy.isa<SizeType>())
       op.emitError() << "if none of the operands can hold error values then "
                         "the result must be of type `index`";
void GetExtentOp::build(OpBuilder &builder, OperationState &result, Value shape,
   }
 }
 
 //===----------------------------------------------------------------------===//
 // RankOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(shape::RankOp op) {
-  Type argTy = op.shape().getType();
-  Type resultTy = op.rank().getType();
-  if (argTy.isa<ShapeType>() && !resultTy.isa<SizeType>())
+  if (op.shape().getType().isa<ShapeType>() &&
+      !op.rank().getType().isa<SizeType>())
     return op.emitOpError()
            << "if operand is of type `shape` then the result must be of type "
               "`size` to propagate potential errors";
   return success();
 }
 
 OpFoldResult shape::RankOp::fold(ArrayRef<Attribute> operands) {
   auto shape = operands[0].dyn_cast_or_null<DenseIntElementsAttr>();
OpFoldResult NumElementsOp::fold(ArrayRef<Attribute> operands) {
   APInt product(64, 1);
   for (auto value : shape.cast<DenseIntElementsAttr>())
     product *= value;
   Builder builder(getContext());
   return builder.getIndexAttr(product.getLimitedValue());
 }
 
 //===----------------------------------------------------------------------===//
+// MulOp
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(MulOp op) {
+  Type resultTy = op.result().getType();
+  if (isErrorPropagationPossible({op.lhs().getType(), op.rhs().getType()})) {
+    if (!resultTy.isa<SizeType>())
+      return op.emitOpError()
+             << "if at least one of the operands can hold error values then "
+                "the result must be of type `size` to propagate them";
+  } else {
+    if (resultTy.isa<SizeType>())
+      return op.emitError() << "if none of the operands can hold error values "
+                               "then the result must be of type `index`";
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // ShapeOfOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult ShapeOfOp::fold(ArrayRef<Attribute>) {
   auto type = getOperand().getType().dyn_cast<ShapedType>();
   if (!type || !type.hasStaticShape())
     return nullptr;
   Builder builder(getContext());
   return builder.getIndexTensorAttr(type.getShape());
 }
 
 static LogicalResult verify(ShapeOfOp op) {
-  Type argTy = op.arg().getType();
   Type resultTy = op.result().getType();
-  if (argTy.isa<ValueShapeType>()) {
+  if (isErrorPropagationPossible(op.arg().getType())) {
     if (!resultTy.isa<ShapeType>())
       return op.emitOpError()
              << "if operand is of type `value_shape` then the result must be "
                 "of type `shape` to propagate potential error shapes";
   } else {
-    assert(argTy.isa<ShapedType>());
     if (resultTy != getExtentTensorType(op.getContext()))
       return op.emitOpError() << "if operand is a shaped type then the result "
                                  "must be an extent tensor";
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/Shape/Transforms/ShapeToShapeLowering.cpp

NumElementsOpConverter::matchAndRewrite(NumElementsOp op,
                                         PatternRewriter &rewriter) const {
   auto loc = op.getLoc();
   Value init = rewriter.create<ConstSizeOp>(loc, rewriter.getIndexAttr(1));
   ReduceOp reduce = rewriter.create<ReduceOp>(loc, op.shape(), init);
 
   // Generate reduce operator.
   Block *body = reduce.getBody();
   OpBuilder b = OpBuilder::atBlockEnd(body);
-  Value product =
-      b.create<MulOp>(loc, body->getArgument(1), body->getArgument(2));
+  Value product = b.create<MulOp>(loc, b.getType<SizeType>(),
+                                  body->getArgument(1), body->getArgument(2));
   b.create<YieldOp>(loc, product);
 
   rewriter.replaceOp(op, reduce.result());
   return success();
 }
 
 namespace {
 struct ShapeToShapeLowering

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

 }
 
 // -----
 
 // Lower binary ops.
 // CHECK-LABEL: @binary_ops
 // CHECK-SAME: (%[[LHS:.*]]: index, %[[RHS:.*]]: index)
 func @binary_ops(%lhs : !shape.size, %rhs : !shape.size) {
+  // CHECK: addi %[[LHS]], %[[RHS]] : index
   %sum = "shape.add"(%lhs, %rhs) : (!shape.size, !shape.size) -> !shape.size
-  // CHECK-NEXT: addi %[[LHS]], %[[RHS]] : index
-  %product = shape.mul %lhs, %rhs
-  // CHECK-NEXT: muli %[[LHS]], %[[RHS]] : index
+  return
+}
+
+// -----
+
+// Lower binary ops.
+// CHECK-LABEL: @binary_ops
+// CHECK-SAME: (%[[LHS:.*]]: index, %[[RHS:.*]]: index)
+func @binary_ops(%lhs : index, %rhs : index) {
+  // CHECK: muli %[[LHS]], %[[RHS]] : index
+  %product = shape.mul %lhs, %rhs : index, index -> index
   return
 }
 
 // -----
 
 // Lower `shape_of` for statically shaped tensor.
 // CHECK-LABEL: @shape_of_stat
 // CHECK-SAME: (%[[ARG:.*]]: tensor<1x2x3xf32>)

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.shape -> !shape.size {
     ^bb0(%index: index, %dim: !shape.size):
       shape.yield %dim : !shape.size
   }
+  return
 }
 
 // -----
 
 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.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
   }
+  return
 }
 
 // -----
 
 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 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
   }
+  return
 }
 
 // -----
 
 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
   }
+  return
 }
 
 // -----
 
 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) : !shape.shape -> f32 {
     ^bb0(%index: index, %dim: !shape.size, %lci: !shape.size):
       shape.yield
   }
+  return
 }
 
 // -----
 
 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.shape -> !shape.size {
     ^bb0(%index: index, %dim: !shape.size, %lci: !shape.size):
       shape.yield %dim, %dim : !shape.size, !shape.size
   }
+  return
 }
 
 // -----
 
 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.shape -> !shape.size {
     ^bb0(%index: index, %dim: !shape.size, %lci: !shape.size):
       %c0 = constant 1 : index
       shape.yield %c0 : index
   }
+  return
 }
 
 // -----
 
 func @assuming_all_op_too_few_operands() {
   // expected-error@+1 {{no operands specified}}
   %w0 = shape.assuming_all
   return
 }
 
 // -----
 
 func @shape_of(%value_arg : !shape.value_shape,
                %shaped_arg : tensor<?x3x4xf32>) {
   // expected-error@+1 {{if operand is of type `value_shape` then the result must be of type `shape` to propagate potential error shapes}}
   %0 = shape.shape_of %value_arg : !shape.value_shape -> tensor<?xindex>
+  return
 }
 
 // -----
 
 func @shape_of(%value_arg : !shape.value_shape,
                %shaped_arg : tensor<?x3x4xf32>) {
   // expected-error@+1 {{if operand is a shaped type then the result must be an extent tensor}}
   %1 = shape.shape_of %shaped_arg : tensor<?x3x4xf32> -> !shape.shape
+  return
 }
 
 // -----
 
 func @rank(%arg : !shape.shape) {
   // expected-error@+1 {{if operand is of type `shape` then the result must be of type `size` to propagate potential errors}}
   %0 = shape.rank %arg : !shape.shape -> index
+  return
 }
 
 // -----
 
 func @get_extent_error_free(%arg : tensor<?xindex>) -> !shape.size {
   %c0 = constant 0 : index
   // expected-error@+1 {{if none of the operands can hold error values then the result must be of type `index`}}
   %result = shape.get_extent %arg, %c0 : tensor<?xindex>, index -> !shape.size
   return %result : !shape.size
 }
 
 // -----
 
 func @get_extent_error_possible(%arg : tensor<?xindex>) -> index {
   %c0 = shape.const_size 0
   // expected-error@+1 {{if at least one of the operands can hold error values then the result must be of type `size` to propagate them}}
   %result = shape.get_extent %arg, %c0 : tensor<?xindex>, !shape.size -> index
   return %result : index
 }
 
+// -----
+
+func @mul_error_free(%arg : index) -> !shape.size {
+  // expected-error@+1 {{if none of the operands can hold error values then the result must be of type `index`}}
+  %result = shape.mul %arg, %arg : index, index -> !shape.size
+  return %result : !shape.size
+}
+
+// -----
+
+func @mul_error_possible(%lhs : !shape.size, %rhs : index) -> index {
+  // expected-error@+1 {{if at least one of the operands can hold error values then the result must be of type `size` to propagate them}}
+  %result = shape.mul %lhs, %rhs : !shape.size, index -> index
+  return %result : index
+}
+

mlir/test/Dialect/Shape/ops.mlir

 // 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 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.size, !shape.size -> !shape.size
       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
+      %acc_next = shape.mul %acc, %extent : index, index -> 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)

 }
 
 func @broadcastable_on_extent_tensors(%lhs : tensor<?xindex>,
                                       %rhs : tensor<?xindex>) {
   %w0 = shape.cstr_broadcastable %lhs, %rhs : tensor<?xindex>, tensor<?xindex>
   return
 }
 
-func @test_mul(%lhs: !shape.size, %rhs: !shape.size) -> !shape.size {
-  %product = shape.mul %lhs, %rhs
-  return %product: !shape.size
+func @mul(%size_arg : !shape.size, %index_arg : index) {
+  %size_prod = shape.mul %size_arg, %size_arg
+      : !shape.size, !shape.size -> !shape.size
+  %index_prod = shape.mul %index_arg, %index_arg : index, index -> index
+  %mixed_prod = shape.mul %size_arg, %index_arg
+      : !shape.size, index -> !shape.size
+  return
 }
 
 func @const_size() {
   // CHECK: %c1 = shape.const_size 1
   // CHECK: %c2 = shape.const_size 2
   // CHECK: %c2_0 = shape.const_size 2
   %0 = shape.const_size 1
   %1 = shape.const_size 2