[mlir] Add a shape op that always returns a successful witness

[tpopp]

Summary

This will later be used during canonicalization steps to replace
statically known passing constraints. This might later be removed once
witnesses can be represented as Attributes.

[tpopp]

Remove named TODO.

[silvas]

LGTM, thanks :)

See my comment in the other patch about using this in ShapeDialect::materializeConstant

(I'm on the fence about having a const_witness {true,false} vs true_witness/false_witness ops; that latter is a bit easier to match with isa<TrueWitnessOp> so I tend to lean towards it, but the former is more consistent with const_shape/const_size and maybe a simple helper can recuperate the advantages of isa<TrueWitnessOp>'s simplicity; feel free to choose whatever makes sense)

[jpienaar]

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

570

I'd be +1 this for consistency. If we have a witness attribute then I think we could just use std.const too.

(then little helper function to check isTrueWitness(Operation* op), which we could add already)

574

Summaries are single line sentence fragments that don't end with punctuation (https://mlir.llvm.org/docs/OpDefinitions/#operation-documentation). I see we are a bit inconsistent in this file on that already.

Might just submit a quick NFC change resolving that,

576

Should we perhaps start with some lead up before the example? And put the MLIR in a code block? (I added syntax highlighting especially for this :-))

This operation produces a constant true witness value. This often arises during constant folding/canonicalization, e.g.,

mlir
  ...

Slightly repetitive with the summary, but given custom asm syntax these will be split in the generated docs.

[herhut]

Thanks!

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

570

I like the isTrueWitness to hide how this is implemented. Then we can move this to ConstantLike later without changing users.

[tpopp]

Instead of true_witness, have const_witness that can be true or false.

[rriddle]

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

577

nit: Remove the indent inside of this block.

mlir/lib/Dialect/Shape/IR/Shape.cpp

45

nit: Remove trivial braces.

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

let description = [{
 
     "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
+      %w0 = shape.cstr_broadcastable [2,2], [3,1,2] // Passing
       %w1 = shape.cstr_broadcastable [2,2], [3,2] // Failure
-      %w2 = shape.cstr_eq [1,2], [1,2], [1,2] // Success
+      %w2 = shape.cstr_eq [1,2], [1,2], [1,2] // Passing
       %wf = shape.assuming_all %w0, %w1 // Failure
-      %wt = shape.assuming_all %w0, %w2 // Success
+      %wt = shape.assuming_all %w0, %w2 // Passing
     ```
   }];
 
   let arguments = (ins Variadic<Shape_WitnessType>:$inputs);
   let results = (outs Shape_WitnessType:$result);
 
   let assemblyFormat = "$inputs attr-dict";
 }
def Shape_CstrBroadcastableOp : Shape_Op<"cstr_broadcastable", []> {
   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
+      %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 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
+      %w0 = shape.cstr_eq [1,2], [1,2], [1,2] // Passing
       %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.
+def Shape_ConstWitnessOp : Shape_Op<"const_witness", [ConstantLike, NoSideEffect]> {
+  let summary = "An operation that returns a statically known witness value";
+  let description = [{
+  This operation represents a statically known witness result. This can be
+  often used to canonicalize/fold constraint and assuming code that will always
+  pass.
+
+  ```mlir
+  %0 = shape.const_shape [1,2,3]
+  %1 = shape.const_shape [1, 2, 3]
+  %w0 = shape.cstr_eq(%0, %1) // Can be folded to "const_witness true"
+  %w1 = shape.const_witness true
+  %w2 = shape.assuming_all(%w0, %w2) // Can be folded to "const_witness true"
+  ```
+  }];
+  let arguments = (ins BoolAttr:$passing);
+  let results = (outs Shape_WitnessType:$result);
+
+  let assemblyFormat = "$passing attr-dict";
+
+  let hasFolder = 1;
+}
+
 
 #endif // SHAPE_OPS

mlir/lib/Dialect/Shape/IR/Shape.cpp

Operation *ShapeDialect::materializeConstant(OpBuilder &builder,
                                              Location loc) {
   if (auto shapeType = type.dyn_cast<ShapeType>()) {
     return builder.create<ConstShapeOp>(loc, type,
                                         value.cast<DenseIntElementsAttr>());
   }
   if (auto sizeType = type.dyn_cast<SizeType>()) {
     return builder.create<ConstSizeOp>(loc, type, value.cast<IntegerAttr>());
   }
+  if (auto witnessType = type.dyn_cast<WitnessType>()) {
+    return builder.create<ConstWitnessOp>(loc, type, value.cast<BoolAttr>());
+  }
   return nullptr;
 }
 
 /// Parse a type registered to this dialect.
 Type ShapeDialect::parseType(DialectAsmParser &parser) const {
   StringRef keyword;
   if (parser.parseKeyword(&keyword))
     return Type();

 
 //===----------------------------------------------------------------------===//
 // ConstSizeOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult ConstSizeOp::fold(ArrayRef<Attribute>) { return valueAttr(); }
 
 //===----------------------------------------------------------------------===//
+// ConstWitnessOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult ConstWitnessOp::fold(ArrayRef<Attribute>) { return passingAttr(); }
+
+//===----------------------------------------------------------------------===//
 // IndexToSizeOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult IndexToSizeOp::fold(ArrayRef<Attribute> operands) {
   // Constant values of both types, `shape.size` and `index`, are represented as
   // `IntegerAttr`s which makes constant folding simple.
   if (Attribute arg = operands[0])
     return arg;

mlir/test/Dialect/Shape/ops.mlir

func @test_shape_of(%arg0: tensor<?xf32>) -> !shape.shape {
   return %0 : !shape.shape
 }
 
 func @test_constraints() {
   %0 = shape.const_shape []
   %1 = shape.const_shape [1, 2, 3]
   %w0 = shape.cstr_broadcastable %0, %1
   %w1 = shape.cstr_eq %0, %1
-  %w3 = shape.assuming_all %w0, %w1
-  shape.assuming %w3 -> !shape.shape {
+  %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 @test_mul(%lhs: !shape.size, %rhs: !shape.size) -> !shape.size {
   %product = shape.mul %lhs, %rhs
   return %product: !shape.size
 }