[MLIR][Shape] Allow `num_elements` to operate on extent tensors

[frgossen]

Summary

Depends On D84444

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

def Shape_MulOp : Shape_Op<"mul", [Commutative, NoSideEffect]> {
   let verifier = [{ return ::verify(*this); }];
   let hasFolder = 1;
 }
 
 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
-    %product = shape.mul %lhs, %rhs
-    ```
+    extents. If the argument is of type `shape` then the result will be of type
+    `size` and potential errors will be propagated. Otherwise, if the argument
+    is and extent tensor `tensor<?xindex>` then the result will be of type
+    `index`.
   }];
 
-  let arguments = (ins Shape_ShapeType:$shape);
-  let results = (outs Shape_SizeType:$result);
+  let arguments = (ins Shape_ShapeOrExtentTensorType:$shape);
+  let results = (outs Shape_SizeOrIndexType:$result);
 
-  let assemblyFormat = "$shape attr-dict";
+  let assemblyFormat = "$shape `:` type($shape) `->` type($result) attr-dict";
 
   let hasFolder = 1;
+  let verifier = [{ return ::verifySizeOrIndexOp(*this); }];
 }
 
 def Shape_ReduceOp : Shape_Op<"reduce",
     [SingleBlockImplicitTerminator<"YieldOp">]> {
   let summary = "Returns an expression reduced over a shape or extent tensor";
   let description = [{
     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

mlir/test/Dialect/Shape/canonicalize.mlir

 // -----
 
 // Fold number of elements computation.
 // CHECK-LABEL: func @num_elements
 func @num_elements() -> !shape.size {
   // CHECK-NOT: shape.const_shape
   %shape = shape.const_shape [4, 5, 6] : !shape.shape
   // CHECK-NOT: shape.num_elements
-  %num_elements = shape.num_elements %shape
+  %num_elements = shape.num_elements %shape : !shape.shape -> !shape.size
   // CHECK: %[[NUM:.*]] = shape.const_size 120
   // CHECK-NEXT: return %[[NUM]] : !shape.size
   return %num_elements : !shape.size
 }
 
 // -----
 
 // No folding.
 // CHECK-LABEL: func @nonfoldable_num_elements
 func @nonfoldable_num_elements(%shape : !shape.shape) -> !shape.size {
   // CHECK-NOT: shape.const_{{.*}}
-  %num_elements = shape.num_elements %shape
+  %num_elements = shape.num_elements %shape : !shape.shape -> !shape.size
   return %num_elements : !shape.size
 }
 
 // -----
 
 // Basic folding.
 // CHECK-LABEL: func @basic
 func @basic() -> index {

mlir/test/Dialect/Shape/invalid.mlir

 // -----
 
 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
 }
 
+// -----
+
+func @num_elements_error_free(%arg : tensor<?xindex>) -> !shape.size {
+  // expected-error@+1 {{if none of the operands can hold error values then the result must be of type `index`}}
+  %result = shape.num_elements %arg : tensor<?xindex> -> !shape.size
+  return %result : !shape.size
+}
+
+// -----
+
+func @num_elements_error_possible(%arg : !shape.shape) -> 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.num_elements %arg : !shape.shape -> index
+  return %result : index
+}
+

mlir/test/Dialect/Shape/ops.mlir

func @any() {
   %1 = shape.const_shape [4, 5, 6] : !shape.shape
   %2 = shape.any %0, %1 : !shape.shape
   %3 = shape.const_shape [1, 2, 3] : tensor<?xindex>
   %4 = shape.const_shape [4, 5, 6] : tensor<?xindex>
   %5 = shape.any %3, %4 : tensor<?xindex>
   return
 }
 
+func @num_elements_extent_tensor(%arg : tensor<?xindex>) -> index {
+  %result = shape.num_elements %arg : tensor<?xindex> -> index
+  return %result : index
+}
+
+func @num_elements_shape(%arg : !shape.shape) -> !shape.size {
+  %result = shape.num_elements %arg : !shape.shape -> !shape.size
+  return %result : !shape.size
+}
+
+

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) -> !shape.size {
+// CHECK-LABEL: func @num_elements_to_reduce
+// CHECK-SAME:  ([[ARG:%.*]]: !shape.shape) -> !shape.size
 func @num_elements_to_reduce(%shape : !shape.shape) -> !shape.size {
-  %num_elements = shape.num_elements %shape
+  %num_elements = shape.num_elements %shape : !shape.shape -> !shape.size
   return %num_elements : !shape.size
 }
 // CHECK: [[C1:%.*]] = shape.const_size 1
 // 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]] : !shape.size
 // CHECK: }
 // CHECK: return [[NUM_ELEMENTS]] : !shape.size