[mlir] Add custom assembly formats to shape.witness ops.

The assembly formats are essentially the generic forms without
quotations and type information.

Differential Revision: https://reviews.llvm.org/D80180

diff --git a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
index 2e521db..7d62ceb 100644 (file)
--- a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
+++ b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
@@ -381,13 +381,15 @@ def Shape_AnyOp : Shape_Op<"any",
 
     Example:
     ```mlir
-      %s0 = shape.any([2,?], [?,3]) // [2,3]
-      %s1 = shape.any([?,?], [1,2]) // [1,2]
+      %s0 = shape.any [2,?], [?,3] // [2,3]
+      %s1 = shape.any [?,?], [1,2] // [1,2]
     ```
   }];
 
   let arguments = (ins Variadic<Shape_ShapeType>:$inputs);
   let results = (outs Shape_ShapeType:$result);
+
+  let assemblyFormat = "$inputs attr-dict";
 }
 
 def Shape_AssumingAllOp : Shape_Op<"assuming_all", [NoSideEffect]> {
@@ -403,16 +405,18 @@ def Shape_AssumingAllOp : Shape_Op<"assuming_all", [NoSideEffect]> {
 
     Example:
     ```mlir
-      %w0 = shape.cstr_broadcastable([2,2], [3,1,2]) // Success
-      %w1 = shape.cstr_broadcastable([2,2], [3,2])    // Failure
-      %w2 = shape.cstr_eq([1,2], [1,2], [1,2]) // Success
-      %wf = shape.assume_all(%w0, %w1) // Failure
-      %wt = shape.assume_all(%w0, %w2) // Success
+      %w0 = shape.cstr_broadcastable [2,2], [3,1,2] // Success
+      %w1 = shape.cstr_broadcastable [2,2], [3,2] // Failure
+      %w2 = shape.cstr_eq [1,2], [1,2], [1,2] // Success
+      %wf = shape.assuming_all %w0, %w1 // Failure
+      %wt = shape.assuming_all %w0, %w2 // Success
     ```
   }];
 
   let arguments = (ins Variadic<Shape_WitnessType>:$inputs);
   let results = (outs Shape_WitnessType:$result);
+
+  let assemblyFormat = "$inputs attr-dict";
 }
 
 def Shape_AssumingOp : Shape_Op<"assuming",
@@ -451,6 +455,8 @@ def Shape_AssumingYieldOp : Shape_Op<"assuming_yield",
     OpBuilder<"OpBuilder &builder, OperationState &result",
     [{ /* nothing to do */ }]>
   ];
+
+  let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
 }
 
 def Shape_CstrBroadcastableOp : Shape_Op<"cstr_broadcastable", []> {
@@ -463,13 +469,15 @@ def Shape_CstrBroadcastableOp : Shape_Op<"cstr_broadcastable", []> {
 
     Example:
     ```mlir
-      %w0 = shape.cstr_broadcastable([2,2], [3,1,2])  // Success
-      %w1 = shape.cstr_broadcastable([2,2], [3,2])    // Failure
+      %w0 = shape.cstr_broadcastable [2,2], [3,1,2] // Success
+      %w1 = shape.cstr_broadcastable [2,2], [3,2] // Failure
     ```
   }];
 
   let arguments = (ins Shape_ShapeType:$lhs, Shape_ShapeType:$rhs);
   let results = (outs Shape_WitnessType:$result);
+
+  let assemblyFormat = "$lhs `,` $rhs attr-dict";
 }
 
 def Shape_CstrEqOp : Shape_Op<"cstr_eq", []> {
@@ -481,12 +489,14 @@ def Shape_CstrEqOp : Shape_Op<"cstr_eq", []> {
 
     Example:
     ```mlir
-      %w0 = shape.cstr_eq([1,2], [1,2], [1,2]) // Success
-      %w1 = shape.cstr_eq([2,2], [1,2])        // Failure
+      %w0 = shape.cstr_eq [1,2], [1,2], [1,2] // Success
+      %w1 = shape.cstr_eq [2,2], [1,2] // Failure
     ```
   }];
   let arguments = (ins Variadic<Shape_ShapeType>:$inputs);
   let results = (outs Shape_WitnessType:$result);
+
+  let assemblyFormat = "$inputs attr-dict";
 }
 
 

diff --git a/mlir/test/Dialect/Shape/ops.mlir b/mlir/test/Dialect/Shape/ops.mlir
index 7934f4e..42eaf58 100644 (file)
--- a/mlir/test/Dialect/Shape/ops.mlir
+++ b/mlir/test/Dialect/Shape/ops.mlir
@@ -71,12 +71,12 @@ func @test_shape_of(%arg0: tensor<?xf32>) -> !shape.shape {
 func @test_constraints() {
   %0 = shape.const_shape []
   %1 = shape.const_shape [1, 2, 3]
-  %w0 = "shape.cstr_broadcastable"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.witness
-  %w1 = "shape.cstr_eq"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.witness
-  %w3 = "shape.assuming_all"(%w0, %w1) : (!shape.witness, !shape.witness) -> !shape.witness
+  %w0 = shape.cstr_broadcastable %0, %1
+  %w1 = shape.cstr_eq %0, %1
+  %w3 = shape.assuming_all %w0, %w1
   shape.assuming %w3 -> !shape.shape {
-    %2 = "shape.any"(%0, %1)  : (!shape.shape, !shape.shape) -> !shape.shape
-    "shape.assuming_yield"(%2) : (!shape.shape) -> ()
+    %2 = shape.any %0, %1
+    shape.assuming_yield %2 : !shape.shape
   }
   return
 }