[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

502

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)

506

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,

508

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

502

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

526

nit: Remove the indent inside of this block.

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

45 ↗

(On Diff #268475)

nit: Remove trivial braces.

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

def Shape_ConcatOp : Shape_Op<"concat",
 
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Shape constraint related ops.
 //===----------------------------------------------------------------------===//
 
-//TODO(tpopp): Move the code below and witnesses to a different file.
+//TODO: Move the code below and witnesses to a different file.
 def Shape_AnyOp : Shape_Op<"any",
     [NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
   let summary = "Return any combination of the input shapes.";
   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
let description = [{
     ```
   }];
   let arguments = (ins Variadic<Shape_ShapeType>:$inputs);
   let results = (outs Shape_WitnessType:$result);
 
   let assemblyFormat = "$inputs attr-dict";
 }
 
+// TODO: Support witness attributes and then make this ConstantLike.
+// Note: This operation might be replaced with a general op that takes a
+// True/False Attribute.
+def Shape_TrueWitnessOp : Shape_Op<"true_witness", [NoSideEffect]> {
+  let summary = "An operation that returns a successful witness.";
+  let description = [{
+      %0 = shape.const_shape [1, 2, 3]
+      %1 = shape.const_shape [1, 2, 3]
+      %w0 = shape.cstr_eq(%0, %1) // Can be canonicalized to true_witness
+      %w1 = shape.true_witness
+      %w2 = shape.assuming_all(%w0, %w2) // Can be canonicalized to true_witness
+  }];
+  let builders = [OpBuilder<
+    "OpBuilder &b, OperationState &result",
+    "build(b, result, ::mlir::shape::WitnessType::get(b.getContext()));"
+  >];
 
-// Canonicalization patterns.
+  let assemblyFormat = "attr-dict";
+  let results = (outs Shape_WitnessType:$result);
+}
 
 #endif // SHAPE_OPS

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
+  %w2 = shape.true_witness
+  %w3 = shape.assuming_all %w0, %w1, %w2
   shape.assuming %w3 -> !shape.shape {
     %2 = shape.any %0, %1
     shape.assuming_yield %2 : !shape.shape
   }
   return
 }