[mlir][shape] Add `shape.cstr_require %bool`

[silvas]

Summary

This op is a catch-all for creating witnesses from various random kinds of constraints. In particular, when dealing with extents directly, which are of index type, one can directly use std ops for calculating the predicates, and then use cstr_require for the final conversion to a
witness.

[silvas]

update

[jpienaar]

Thanks, I am tempted to suggest we add "use this only if more specific form is not available" (e.g., use this for cases not expressible with the others). So that we push to simpler forms.

[silvas]

+1, added some verbiage to that effect

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

def Shape_ConstWitnessOp : Shape_Op<"const_witness", [ConstantLike, NoSideEffect]> {
   let arguments = (ins BoolAttr:$passing);
   let results = (outs Shape_WitnessType:$result);
 
   let assemblyFormat = "$passing attr-dict";
 
   let hasFolder = 1;
 }
 
+def Shape_CstrRequireOp : Shape_Op<"cstr_require", []> {
+  let summary = "Represents a runtime assertion that an i1 is `true`";
+  let description = [{
+    Represents a runtime assretion that an i1 is true. It returns a
+    !shape.witness to order this assertion.
+
+    For simplicity, prefer using other cstr_* ops if they are available for a
+    given constraint.
+
+    Example:
+    ```mlir
+    %bool = ...
+    %w0 = shape.cstr_require %bool // Passing if `%bool` is true.
+    ```
+  }];
+  let arguments = (ins I1:$pred);
+  let results = (outs Shape_WitnessType:$result);
+
+  let assemblyFormat = "$pred attr-dict";
+
+  let hasFolder = 1;
+}
 
 #endif // SHAPE_OPS

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

 
 //===----------------------------------------------------------------------===//
 // ConstWitnessOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult ConstWitnessOp::fold(ArrayRef<Attribute>) { return passingAttr(); }
 
 //===----------------------------------------------------------------------===//
+// CstrRequireOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult CstrRequireOp::fold(ArrayRef<Attribute> operands) {
+  return operands[0];
+}
+
+//===----------------------------------------------------------------------===//
 // ShapeEqOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult ShapeEqOp::fold(ArrayRef<Attribute> operands) {
   auto lhs = operands[0].dyn_cast_or_null<DenseIntElementsAttr>();
   if (lhs == nullptr)
     return {};
   auto rhs = operands[1].dyn_cast_or_null<DenseIntElementsAttr>();

mlir/test/Dialect/Shape/canonicalize.mlir

func @f(%arg0: !shape.shape, %arg1: !shape.shape) {
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %0 = shape.cstr_eq %arg0, %arg1
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
+// cstr_require with constant can be folded
+// CHECK-LABEL: func @cstr_require_fold
+func @cstr_require_fold() {
+  // CHECK-NEXT: shape.const_witness true
+  // CHECK-NEXT: consume.witness
+  // CHECK-NEXT: return
+  %true = constant true
+  %0 = shape.cstr_require %true
+  "consume.witness"(%0) : (!shape.witness) -> ()
+  return
+}
+
+// -----
+// cstr_require without constant cannot be folded
+// CHECK-LABEL: func @cstr_require_no_fold
+func @cstr_require_no_fold(%arg0: i1) {
+  // CHECK-NEXT: shape.cstr_require
+  // CHECK-NEXT: consume.witness
+  // CHECK-NEXT: return
+  %0 = shape.cstr_require %arg0
+  "consume.witness"(%0) : (!shape.witness) -> ()
+  return
+}
 
+// -----
 // assuming_all with known passing witnesses can be folded
 // CHECK-LABEL: func @f
 func @f() {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %0 = shape.const_witness true
   %1 = shape.const_witness true

mlir/test/Dialect/Shape/ops.mlir

 func @test_shape_of(%arg0: tensor<?xf32>) -> tensor<?xindex> {
   %0 = shape.shape_of %arg0 : tensor<?xf32> -> tensor<?xindex>
   return %0 : tensor<?xindex>
 }
 
 func @test_constraints() {
   %0 = shape.const_shape [] : !shape.shape
   %1 = shape.const_shape [1, 2, 3] : !shape.shape
+  %true = constant true
   %w0 = shape.cstr_broadcastable %0, %1 : !shape.shape, !shape.shape
   %w1 = shape.cstr_eq %0, %1
   %w2 = shape.const_witness true
   %w3 = shape.const_witness false
-  %w4 = shape.assuming_all %w0, %w1, %w2, %w3
-  shape.assuming %w4 -> !shape.shape {
+  %w4 = shape.cstr_require %true
+  %w_all = shape.assuming_all %w0, %w1, %w2, %w3, %w4
+  shape.assuming %w_all -> !shape.shape {
     %2 = "shape.any"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
     shape.assuming_yield %2 : !shape.shape
   }
   return
 }
 
 func @broadcastable_on_extent_tensors(%lhs : tensor<?xindex>,
                                       %rhs : tensor<?xindex>) {