[MLIR][Shape] Add custom assembly format for `shape.any`

[frgossen]

Summary

Add custom assembly format for shape.any with variadic operands.

[pifon2a]

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

584

can it be done with assemblyFormat?

[frgossen]

Use assembly format

[frgossen]

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

584

Did not expect that assemblyFormat could handle variadic inputs. Thanks!

[herhut]

Thanks for cleaning this up!

[jpienaar]

Thanks!

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

let description = [{
     %s0 = shape.any [2,?], [?,3] // [2,3]
     %s1 = shape.any [?,?], [1,2] // [1,2]
     ```
   }];
 
   let arguments = (ins Variadic<Shape_ShapeOrExtentTensorType>:$inputs);
   let results = (outs Shape_ShapeOrExtentTensorType:$result);
 
+  let assemblyFormat = "$inputs attr-dict `:` type($inputs) `->` type($result)";
+
   let hasFolder = 1;
 }
 
 def Shape_AssumingAllOp : Shape_Op<"assuming_all", [Commutative, NoSideEffect]> {
   let summary = "Return a logical AND of all witnesses";
   let description = [{
     Used to simplify constraints as any single failing precondition is enough
     to prevent execution.

mlir/test/Dialect/Shape/canonicalize.mlir

 // -----
 
 // any can be replaced with a constant input if it has one.
 // CHECK-LABEL: func @f
 func @f(%arg : !shape.shape) -> !shape.shape {
   // CHECK-NEXT: %[[CS:.*]] = shape.const_shape
   // CHECK-NEXT: return %[[CS]]
   %0 = shape.const_shape [2, 3, 4] : !shape.shape
-  %1 = "shape.any"(%0, %arg) : (!shape.shape, !shape.shape) -> !shape.shape
+  %1 = shape.any %0, %arg : !shape.shape, !shape.shape -> !shape.shape
   return %1 : !shape.shape
 }
 
 // -----
 
 // any can be replaced with a constant input if it has one.
 // CHECK-LABEL: func @f
 func @f(%arg : tensor<?xindex>) -> tensor<?xindex> {
   // CHECK-NEXT: %[[CS:.*]] = shape.const_shape [2, 3, 4] : tensor<?xindex>
   // CHECK-NEXT: return %[[CS]] : tensor<?xindex>
   %0 = shape.const_shape [2, 3, 4] : tensor<?xindex>
-  %1 = "shape.any"(%0, %arg) : (tensor<?xindex>, tensor<?xindex>) -> tensor<?xindex>
+  %1 = shape.any %0, %arg : tensor<?xindex>, tensor<?xindex> -> tensor<?xindex>
   return %1 : tensor<?xindex>
 }
 
 // -----
 
 // Folding of any with partially constant operands is not yet implemented.
 // CHECK-LABEL: func @f
 func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> !shape.shape {
-  // CHECK-NEXT: %[[CS:.*]] = "shape.any"
+  // CHECK-NEXT: %[[CS:.*]] = shape.any
   // CHECK-NEXT: return %[[CS]]
-  %1 = "shape.any"(%arg0, %arg1) : (!shape.shape, !shape.shape) -> !shape.shape
+  %1 = shape.any %arg0, %arg1 : !shape.shape, !shape.shape -> !shape.shape
   return %1 : !shape.shape
 }
 
 // -----
 
 // assuming with a known passing witness can be removed
 // CHECK-LABEL: func @f
 func @f() {

mlir/test/Dialect/Shape/ops.mlir

func @shape_equal_shapes(%a : !shape.value_shape, %b : !shape.value_shape) -> !shape.shape {
   return %2 : !shape.shape
 }
 func @shape_with_shape(%a : !shape.value_shape, %b : !shape.value_shape) -> !shape.shape {
   %0 = shape.shape_of %a : !shape.value_shape -> !shape.shape
   %1 = shape.with_shape %b, %0 : !shape.value_shape, !shape.shape
   %2 = call @shape_equal_shapes(%a, %1) : (!shape.value_shape, !shape.value_shape) -> !shape.shape
   return %2 : !shape.shape
 }
+
+func @any_on_shape(%a : !shape.shape, %b : !shape.shape, %c : !shape.shape)
+    -> !shape.shape {
+  %result = shape.any %a, %b, %c
+      : !shape.shape, !shape.shape, !shape.shape -> !shape.shape
+  return %result : !shape.shape
+}
+
+func @any_on_mixed(%a : tensor<?xindex>,
+                   %b : tensor<?xindex>,
+                   %c : !shape.shape) -> !shape.shape {
+  %result = shape.any %a, %b, %c
+      : tensor<?xindex>, tensor<?xindex>, !shape.shape -> !shape.shape
+  return %result : !shape.shape
+}
+
+func @any_on_extent_tensors(%a : tensor<?xindex>,
+                            %b : tensor<?xindex>,
+                            %c : tensor<?xindex>) -> tensor<?xindex> {
+  %result = shape.any %a, %b, %c
+      : tensor<?xindex>, tensor<?xindex>, tensor<?xindex> -> tensor<?xindex>
+  return %result : tensor<?xindex>
+}