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

[tpopp]

Summary

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

[mehdi_amini]

Thanks!

[herhut]

Reads much nicer!

[silvas]

mlir/test/Dialect/Shape/ops.mlir

74–76

can we remove the parentheses? It's easier to type and read IMO and is consistent with many std and other ops (like addi)

[tpopp]

Remove parantheses

[tpopp]

Match example code in documentation to the actual custom form.

[tpopp]

@silvas I was going to wait because you hadn't accepted the commit, but because (1) I am not working the next 4 days and (2) have other commits using these operations, I'm going to submit this, but can follow up on any other feedback.

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

let description = [{
     This operation takes multiple input shapes and returns some combination of
     their dimensions. This can be best seen with examples below.
 
     The result is undefined, but still side-effect free, in cases where the
     inputs have differing ranks or differ in extents of shared dimensions.
 
     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]> {
   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.
 
     "assuming" operations represent an execution order restriction to the
     compiler, information for dependent code to rely on (by assuming), and
     nothing else. They should not exist after a program is fully lowered and
     ready to execute.
 
     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",
                            [SingleBlockImplicitTerminator<"AssumingYieldOp">,
                             RecursiveSideEffects]> {
   let summary = "Execute the region.";
   let description = [{
     Executes the region assuming all witnesses are true.
def Shape_AssumingYieldOp : Shape_Op<"assuming_yield",
   }];
 
   let arguments = (ins Variadic<AnyType>:$operands);
 
   let builders = [
     OpBuilder<"OpBuilder &builder, OperationState &result",
     [{ /* nothing to do */ }]>
   ];
+
+  let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
 }
 
 def Shape_CstrBroadcastableOp : Shape_Op<"cstr_broadcastable", []> {
   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.
 
     "cstr" operations represent runtime assertions.
 
     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", []> {
   let summary = "Determines if all input shapes are equal.";
   let description = [{
     Given 1 or more input shapes, determine if all shapes are the exact same.
 
     "cstr" operations represent runtime assertions.
 
     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";
 }
 
 
 // Canonicalization patterns.
 
 #endif // SHAPE_OPS

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) : (!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
 }