[mlir] Canonicalization of shape.assuming_all

[tpopp]

Summary

This allows assuming_all to be replaced when all inputs are known to be
statically passing witnesses.

[silvas]

LGTM, but let's look into using a fold for this (the good thing is that all the test cases will carry over :) )

[jpienaar]

mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td

1 ↗

(On Diff #265270)

Missing file header?

4 ↗

(On Diff #265270)

Could we indent by 2 instead?

def AllInputsTrueWitnesses : Constraint<CPred<[
  llvm::all_of($0.getOwner()->getOperands(), [](mlir::Value op) {
    return op.getDefiningOp<mlir::shape::TrueWitnessOp>();
  })}]>>;

12 ↗

(On Diff #265270)

Could we just indent these by 2 instead? [yes we need a formatter ... and many other variants as below]

def ConstantAssumingAll : Pat<
  (Shape_AssumingAllOp:$op $input), (Shape_TrueWitnessOp),
  [(AllInputsTrueWitnesses $op)]>;

[herhut]

Feel free to land. Comments are just suggestions and can also be addressed in a follow up.

mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td

6 ↗

(On Diff #265270)

This would be a great place to use the isTrueWitness helper.

11 ↗

(On Diff #265270)

This could also be formulated as a pattern that drops constant true witnesses from the operands and a pattern that replaces an assuming_all with no operands by a constant true witness. That way we also get removal of constant operands (which is more a cosmetic cleanup and does not really propagate any knowledge).

[tpopp]

Instead of canonicalization, use folding.

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

def Shape_AnyOp : Shape_Op<"any", [NoSideEffect]> {
   }];
 
   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]> {
+def Shape_AssumingAllOp : Shape_Op<"assuming_all", [Commutative, 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] // Passing
       %w1 = shape.cstr_broadcastable [2,2], [3,2] // Failure
       %w2 = shape.cstr_eq [1,2], [1,2], [1,2] // Passing
       %wf = shape.assuming_all %w0, %w1 // Failure
       %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";
+
+  let hasFolder = 1;
 }
 
 def Shape_AssumingOp : Shape_Op<"assuming",
                            [SingleBlockImplicitTerminator<"AssumingYieldOp">,
                             RecursiveSideEffects]> {
   let summary = "Execute the region";
   let description = [{
     Executes the region assuming all witnesses are true.

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

static void print(OpAsmPrinter &p, AssumingOp op) {
   }
   p.printRegion(op.doRegion(),
                 /*printEntryBlockArgs=*/false,
                 /*printBlockTerminators=*/yieldsResults);
   p.printOptionalAttrDict(op.getAttrs());
 }
 
 //===----------------------------------------------------------------------===//
+// AssumingAllOp
+//===----------------------------------------------------------------------===//
+OpFoldResult AssumingAllOp::fold(ArrayRef<Attribute> operands) {
+  // Iterate in reverse to first handle all constant operands. They are
+  // guaranteed to be the tail of the inputs because this is commutative.
+  for (int idx = operands.size() - 1; idx >= 0; idx--) {
+    Attribute a = operands[idx];
+    // Cannot fold if any inputs are not constant;
+    if (!a)
+      return nullptr;
+
+    // We do not need to keep statically known values after handling them in
+    // this method.
+    getOperation()->eraseOperand(idx);
+
+    // Always false if any input is statically known false
+    if (!a.cast<BoolAttr>().getValue())
+      return a;
+  }
+  // If this is reached, all inputs were statically known passing.
+  return BoolAttr::get(true, getContext());
+}
+
+//===----------------------------------------------------------------------===//
 // BroadcastOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult BroadcastOp::fold(ArrayRef<Attribute> operands) {
   if (!operands[0] || !operands[1])
     return nullptr;
   auto lhsShape = llvm::to_vector<6>(
       operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());

mlir/test/Dialect/Shape/canonicalize.mlir

-// RUN: mlir-opt -split-input-file -canonicalize <%s | FileCheck %s --dump-input=fail
+// RUN: mlir-opt -split-input-file -allow-unregistered-dialect -canonicalize <%s | FileCheck %s --dump-input=fail
 
 // -----
 // CHECK-LABEL: func @f
 func @f(%arg0: tensor<2x3x4xf32>) -> !shape.shape {
   // CHECK: shape.const_shape [2, 3, 4]
   %0 = "shape.shape_of"(%arg0) : (tensor<2x3x4xf32>) -> !shape.shape
   return %0 : !shape.shape
 }

 
 // Should not fold.
 // CHECK-LABEL: func @not_const
 func @not_const(%arg0: !shape.shape) -> !shape.size {
   // CHECK: shape.get_extent
   %0 = shape.get_extent %arg0, 3
   return %0 : !shape.size
 }
+
+// -----
+// 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
+  %2 = shape.const_witness true
+  %3 = shape.assuming_all %0, %1, %2
+  "consume.witness"(%3) : (!shape.witness) -> ()
+  return
+}
+
+// -----
+// assuming_all should not be removed if not all witnesses are statically passing.
+//
+// Additionally check that the attribute is moved to the end as this op is
+// commutative.
+// CHECK-LABEL: func @f
+func @f() {
+  // CHECK-NEXT: %[[UNKNOWN:.*]] = "test.source"
+  // CHECK-NEXT: shape.assuming_all %[[UNKNOWN]]
+  // CHECK-NEXT: consume.witness
+  // CHECK-NEXT: return
+  %0 = shape.const_witness true
+  %1 = "test.source"() : () -> !shape.witness
+  %2 = shape.assuming_all %0, %1
+  "consume.witness"(%2) : (!shape.witness) -> ()
+  return
+}