[MLIR][Shape] Allow `cstr_broadcastable` to accept extent tensors

[herhut]

I wonder whether we need to duplicate all the tests to also cover tensor types. Under the hood it would be the same code-path anyway, though.

[frgossen]

Add test cases

[frgossen]

Added two test cases to have some test that the extent tensor variant is also considered. Duplication all tests would

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

def Shape_AssumingYieldOp : Shape_Op<"assuming_yield",
   ];
 
   let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
 }
 
 def Shape_CstrBroadcastableOp : Shape_Op<"cstr_broadcastable", [Commutative]> {
   let summary = "Determines if 2 shapes can be successfully broadcasted";
   let description = [{
-    Given 2 input shapes, return a witness specifying if they are broadcastable.
-    This broadcastable follows the same logic as what shape.broadcast documents.
+    Given two input shapes or extent tensors, return a witness specifying if
+    they are broadcastable. This broadcastable follows the same logic as what
+    shape.broadcast documents.
 
     "cstr" operations represent runtime assertions.
 
     Example:
     ```mlir
     %w0 = shape.cstr_broadcastable [2,2], [3,1,2] // Passing
     %w1 = shape.cstr_broadcastable [2,2], [3,2] // Failure
     ```
   }];
 
-  let arguments = (ins Shape_ShapeType:$lhs, Shape_ShapeType:$rhs);
+  let arguments = (ins Shape_ShapeOrExtentTensorType:$lhs,
+                       Shape_ShapeOrExtentTensorType:$rhs);
   let results = (outs Shape_WitnessType:$result);
 
-  let assemblyFormat = "$lhs `,` $rhs attr-dict";
+  let assemblyFormat = "$lhs `,` $rhs `:` type($lhs) `,` type($rhs) attr-dict";
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 def Shape_CstrEqOp : Shape_Op<"cstr_eq", [Commutative]> {
   let summary = "Determines if all input shapes are equal";
   let description = [{

mlir/test/Dialect/Shape/canonicalize.mlir

 // Broadcastable with broadcastable constant shapes can be removed.
 // CHECK-LABEL: func @f
 func @f() {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %cs0 = shape.const_shape [3, 1]
   %cs1 = shape.const_shape [1, 5]
-  %0 = shape.cstr_broadcastable %cs0, %cs1
+  %0 = shape.cstr_broadcastable %cs0, %cs1 : !shape.shape, !shape.shape
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // Broadcastable with non-broadcastable constant shapes is always false
 // CHECK-LABEL: func @static_non_broadcastable
 func @static_non_broadcastable() {
   // CHECK-NEXT: shape.const_shape
   // CHECK-NEXT: shape.const_shape
   // CHECK-NEXT: shape.cstr_broadcastable
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %cs0 = shape.const_shape [1, 3]
   %cs1 = shape.const_shape [1, 5]
-  %0 = shape.cstr_broadcastable %cs0, %cs1
+  %0 = shape.cstr_broadcastable %cs0, %cs1 : !shape.shape, !shape.shape
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // Broadcastable without guaranteed broadcastable shapes cannot be removed.
 // CHECK-LABEL: func @f
 func @f(%arg0 : !shape.shape) {
   // CHECK-NEXT: shape.const_shape
   // CHECK-NEXT: shape.cstr_broadcastable
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %cs0 = shape.const_shape [1,3]
-  %0 = shape.cstr_broadcastable %arg0, %cs0
+  %0 = shape.cstr_broadcastable %arg0, %cs0 : !shape.shape, !shape.shape
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // Broadcastable with non-constant but known equal shapes can be removed.
 // CHECK-LABEL: func @f
 func @f(%arg0 : !shape.shape) {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
-  %0 = shape.cstr_broadcastable %arg0, %arg0
+  %0 = shape.cstr_broadcastable %arg0, %arg0 : !shape.shape, !shape.shape
+  "consume.witness"(%0) : (!shape.witness) -> ()
+  return
+}
+
+// -----
+
+// Broadcastable canonicalization also works on extent tensors.
+// CHECK-LABEL: func @broadcastable_on_extent_tensors
+func @broadcastable_on_extent_tensors(%arg : tensor<?xindex>) {
+  // CHECK-NEXT: shape.const_witness true
+  // CHECK-NEXT: consume.witness
+  // CHECK-NEXT: return
+  %0 = shape.cstr_broadcastable %arg, %arg : tensor<?xindex>, tensor<?xindex>
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 
 // Fold `rank` based on constant shape.
 // CHECK-LABEL: @fold_rank

 // Canonicalize scalar cstr_broadcastable checks
 // CHECK-LABEL: @cstr_broadcastable_scalar
 func @cstr_broadcastable_scalar(%arg0 : tensor<?xf32>) {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %0 = shape.const_shape []
   %1 = shape.shape_of %arg0 : tensor<?xf32>
-  %2 = shape.cstr_broadcastable %0, %1
+  %2 = shape.cstr_broadcastable %0, %1 : !shape.shape, !shape.shape
   "consume.witness"(%2) : (!shape.witness) -> ()
   return
 }
 
 // -----
 
 // Do not canonicalize cstr_broadcastable checks with 2 unknowns
 // CHECK-LABEL: @cstr_broadcastable_unknown
 func @cstr_broadcastable_unknown(%arg0 : tensor<?xf32>, %arg1 : tensor<?xf32>) {
   // CHECK-NEXT: shape.shape_of %arg0
   // CHECK-NEXT: shape.shape_of %arg1
   // CHECK-NEXT: shape.cstr_broadcastable
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %0 = shape.shape_of %arg0 : tensor<?xf32>
   %1 = shape.shape_of %arg1 : tensor<?xf32>
-  %2 = shape.cstr_broadcastable %0, %1
+  %2 = shape.cstr_broadcastable %0, %1 : !shape.shape, !shape.shape
   "consume.witness"(%2) : (!shape.witness) -> ()
   return
 }
 
 // -----
 
 // Scalars are safe to broadcast to unranked sizes.
 // CHECK-LABEL: @cstr_broadcastable_scalar_unranked
 func @cstr_broadcastable_scalar_unranked(%arg0 : tensor<*xf32>, %arg1 : tensor<index>) {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %0 = shape.shape_of %arg1 : tensor<index>
   %1 = shape.shape_of %arg0 : tensor<*xf32>
-  %2 = shape.cstr_broadcastable %0, %1
+  %2 = shape.cstr_broadcastable %0, %1 : !shape.shape, !shape.shape
   "consume.witness"(%2) : (!shape.witness) -> ()
   return
 }
 
 // -----
 
 // Fold `shape_eq` for equal and constant shapes.
 // CHECK-LABEL: @shape_eq_fold_1

mlir/test/Dialect/Shape/ops.mlir

 func @test_shape_of(%arg0: tensor<?xf32>) -> !shape.shape {
   %0 = shape.shape_of %arg0 : tensor<?xf32>
   return %0 : !shape.shape
 }
 
 func @test_constraints() {
   %0 = shape.const_shape []
   %1 = shape.const_shape [1, 2, 3]
-  %w0 = shape.cstr_broadcastable %0, %1
+  %w0 = shape.cstr_broadcastable %0, %1 : !shape.shape, !shape.shape
   %w1 = shape.cstr_eq %0, %1
   %w2 = shape.const_witness true
   %w3 = shape.const_witness false
   %w4 = shape.assuming_all %w0, %w1, %w2, %w3
   shape.assuming %w4 -> !shape.shape {
     %2 = shape.any %0, %1
     shape.assuming_yield %2 : !shape.shape
   }
   return
 }
 
+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 @const_size() {
   // CHECK: %c1 = shape.const_size 1
   // CHECK: %c2 = shape.const_size 2

mlir/test/Dialect/Shape/remove-shape-constraints.mlir

 // RUN: mlir-opt -allow-unregistered-dialect -split-input-file -remove-shape-constraints -canonicalize <%s | FileCheck %s --dump-input=fail --check-prefixes=CANON,CHECK-BOTH
 // RUN: mlir-opt -allow-unregistered-dialect -split-input-file -remove-shape-constraints <%s | FileCheck %s --dump-input=fail --check-prefixes=REPLACE,CHECK-BOTH
 
 // -----
 // Check that cstr_broadcastable is removed.
 //
 // CHECK-BOTH: func @f
 func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> index {
   // REPLACE-NEXT: %[[WITNESS:.+]] = shape.const_witness true
   // REPLACE-NOT: shape.cstr_eq
   // REPLACE: shape.assuming %[[WITNESS]]
   // CANON-NEXT: test.source
   // CANON-NEXT: return
-  %0 = shape.cstr_broadcastable %arg0, %arg1
+  %0 = shape.cstr_broadcastable %arg0, %arg1 : !shape.shape, !shape.shape
   %1 = shape.assuming %0 -> index {
     %2 = "test.source"() : () -> (index)
     shape.assuming_yield %2 : index
   }
   return %1 : index
 }
 
 // -----

 // -----
 // With a non-const value, we cannot fold away the code, but all constraints
 // should be removed still.
 //
 // CHECK-BOTH: func @f
 func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> index {
   // CANON-NEXT: test.source
   // CANON-NEXT: return
-  %0 = shape.cstr_broadcastable %arg0, %arg1
+  %0 = shape.cstr_broadcastable %arg0, %arg1 : !shape.shape, !shape.shape
   %1 = shape.cstr_eq %arg0, %arg1
   %2 = shape.assuming_all %0, %1
   %3 = shape.assuming %0 -> index {
     %4 = "test.source"() : () -> (index)
     shape.assuming_yield %4 : index
   }
   return %3 : index
 }