[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

Missing file header?

4

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

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

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

11

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",
   }];
 
   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] // 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";
+
+  let hasCanonicalizer = 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/CMakeLists.txt

+set(LLVM_TARGET_DEFINITIONS IR/ShapeCanonicalization.td)
+mlir_tablegen(IR/ShapeCanonicalization.inc -gen-rewriters)
+add_public_tablegen_target(MLIRShapeCanonicalizationIncGen)
+
 add_mlir_dialect_library(MLIRShape
   IR/Shape.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Shape
 
   DEPENDS
   MLIRShapeOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRControlFlowInterfaces
   MLIRDialect
   MLIRInferTypeOpInterface
   MLIRIR
   MLIRSideEffectInterfaces
   )

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

 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/StandardTypes.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 using namespace mlir::shape;
 
+namespace {
+#include "IR/ShapeCanonicalization.inc"
+}
 ShapeDialect::ShapeDialect(MLIRContext *context)
     : Dialect(getDialectNamespace(), context) {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Shape/IR/ShapeOps.cpp.inc"
       >();
   addTypes<ComponentType, ElementType, ShapeType, SizeType, ValueShapeType,
            WitnessType>();
static void print(OpAsmPrinter &p, AssumingOp op) {
   }
   p.printRegion(op.doRegion(),
                 /*printEntryBlockArgs=*/false,
                 /*printBlockTerminators=*/yieldsResults);
   p.printOptionalAttrDict(op.getAttrs());
 }
 
 //===----------------------------------------------------------------------===//
+// AssumingAllOp
+//===----------------------------------------------------------------------===//
+void AssumingAllOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
+  // If all true, replace with true. Folding?
+  patterns.insert<ConstantAssumingAll>(context);
+}
+
+//===----------------------------------------------------------------------===//
 // BroadcastOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 BroadcastOp::inferReturnTypes(MLIRContext *context, Optional<Location> location,
                               ValueRange operands, DictionaryAttr attributes,
                               RegionRange regions,
                               SmallVectorImpl<Type> &inferredReturnTypes) {

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

+include "mlir/Dialect/Shape/IR/ShapeOps.td"
+
+// Constraints
+def AllInputsTrueWitnesses : Constraint<CPred<[{
+        llvm::all_of($0.getOwner()->getOperands(), [](mlir::Value op) {
+          return op.getDefiningOp<mlir::shape::TrueWitnessOp>();
+        })
+      }]>>;
+
+// Canonicalization patterns.
+def ConstantAssumingAll : Pat<(Shape_AssumingAllOp:$op $input),
+          (Shape_TrueWitnessOp),
+          [(AllInputsTrueWitnesses $op)] >;

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
 }

 
 // CHECK-LABEL: func @no_fold
 func @no_fold(%arg0: index) -> !shape.shape {
   // CHECK-NOT: shape.const_shape
   %e0 = constant 3 : index
   %ret = shape.from_extents %e0, %arg0
   return %ret : !shape.shape
 }
+
+// -----
+// assuming_all with known true witnesses can be folded
+// CHECK-LABEL: func @f
+func @f() {
+  // CHECK-NEXT: shape.true_witness
+  // CHECK-NEXT: consume.witness
+  // CHECK-NEXT: return
+  %0 = shape.true_witness
+  %1 = shape.true_witness
+  %2 = shape.true_witness
+  %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 true.
+// CHECK-LABEL: func @f
+func @f() {
+  // CHECK-NEXT: %[[TRUE:.*]] = shape.true_witness
+  // CHECK-NEXT: %[[UNKNOWN:.*]] = "test.source"
+  // CHECK-NEXT: shape.assuming_all %[[TRUE]], %[[UNKNOWN]]
+  // CHECK-NEXT: consume.witness
+  // CHECK-NEXT: return
+  %0 = shape.true_witness
+  %1 = "test.source"() : () -> !shape.witness
+  %2 = shape.assuming_all %0, %1
+  "consume.witness"(%2) : (!shape.witness) -> ()
+  return
+}