[mlir] Add canonicalization for Cstr and Assuming Shape Ops.

[tpopp]

Summary

As a first step, represent statically known passing constraint as a
true op and remove assuming ops that rely on only true_witness ops.

[herhut]

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

483

No TODO with names please.

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

141

Why this change?

194

Why initPosition?

mlir/test/Dialect/Shape/canonicalize.mlir

91

Some tests where folding does not happen would also be great.

[silvas]

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

106

Use a fold.

186

We might not want this pattern always. For example, if we have a pass that does code motion on islands, this canonicalization can prematurely remove the island which will make that pass unable to operate properly. Let's hold off on adding it until we have a compelling use case.

325

This will segfault if getDefiningOp returns null.

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

19

I don't think that TableGen is buying us much here. I prefer loose C++ patterns since they are easier to debug and understand.

mlir/test/Dialect/Shape/canonicalize.mlir

119

Can you make this test have less op surface area? Maybe use an unknown op to capture the witness (or return the witness) instead of having the shape.assuming?

[mehdi_amini]

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

19

Easier to understand is arguable: this can be a matter of getting used to it, easier to debug is an infra issue which should be fixed if any.

C++ pattern aren't ideal because we can't manipulate them statically: even if DRR is far from perfect, at least it is declarative.
In particular with the interpreted rewrite engine coming, looking into authoring such transformation in a non-C++ way is valuable (whether we use DRR or we come up with another authoring language to express these)

[tpopp]

I split this commit up into smaller pieces.

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

141

Accidental

186

I don't find this a compelling argument. Not all code will be in islands and these islands are for the purpose of making code conditional on some assertion. If it's statically known that the assertion is not a problem and the island cannot be removed, then the island is being misused for another purpose and blocking possible compiler optimizations.

194

Changed to opPosition as the rewriter starts where the op is.

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

19

I'm going to leave as tablegen. I'm not a big fan of almost all tablegen patterns for LLVM but am going along with what the infrastructure implies it wants.

mlir/test/Dialect/Shape/canonicalize.mlir

119

I split these into one for each canonicalization.

[silvas]

Looks like this has been shraded into a couple other patches. Resigning from this master one as I reviewed the others.

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

186

Good point!

[herhut]

Please abandon this revision.

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

```mlir
       %s1 = shape.any([?,?], [1,2]) // [1,2]
     ```
   }];
 
   let arguments = (ins Variadic<Shape_ShapeType>:$inputs);
   let results = (outs Shape_ShapeType:$result);
 
   let assemblyFormat = "`(` $inputs `)` attr-dict";
+  let hasCanonicalizer = 1;
 }
 
 def Shape_AssumingAllOp : Shape_Op<"assuming_all", [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.
 
```mlir
       %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.
 
     "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.
   }];
   let arguments = (ins Shape_WitnessType:$condition);
   let regions = (region SizedRegion<1>:$doRegion);
   let results = (outs Variadic<AnyType>:$results);
 
   let printer = [{ return ::print(p, *this); }];
   let parser = [{ return ::parse$cppClass(parser, result); }];
+
+  let hasCanonicalizer = 1;
 }
 
 def Shape_AssumingYieldOp : Shape_Op<"assuming_yield",
                                      [NoSideEffect, ReturnLike, Terminator]> {
   let summary = "Yield operation";
   let description = [{
     This yield operation represents a return operation within the assert_and_exec
     region. The operation takes variable number of operands and produces no
```mlir
       %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";
+  let hasCanonicalizer = 1;
 }
 
 def Shape_CstrEqOp : Shape_Op<"cstr_eq", [NoSideEffect]> {
   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
       %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";
+  let hasCanonicalizer = 1;
 }
 
+// TODO(tpopp): 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/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>();
default:
     llvm_unreachable("unexpected 'shape' type kind");
   }
 }
 
 //===----------------------------------------------------------------------===//
 // AnyOp
 //===----------------------------------------------------------------------===//
 
+namespace {
+// Removes AnyOp with constant shape input.
+// TODO(tpopp): This case can be replaced with folding.
+// TODO(tpopp): Another pattern should be implemented for shapes that can be
+// determined through mixtures of the known dimensions of the inputs.
+struct AnyWithConstant : public OpRewritePattern<AnyOp> {
+  using OpRewritePattern<AnyOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AnyOp op,
+                                PatternRewriter &rewriter) const override {
+    for (auto shapeOp : op.getOperands()) {
+      if (shapeOp.getDefiningOp<ConstShapeOp>()) {
+        rewriter.replaceOp(op, shapeOp);
+        return success();
+      }
+    }
+    return failure();
+  }
+};
+}; // namespace
+
+void AnyOp::getCanonicalizationPatterns(OwningRewritePatternList &patterns,
+                                        MLIRContext *context) {
+  // Replace with constant shape? Maybe folding?
+  patterns.insert<AnyWithConstant>(context);
+}
+
 LogicalResult
 AnyOp::inferReturnTypes(MLIRContext *context, Optional<Location> location,
                         ValueRange operands, DictionaryAttr attributes,
                         RegionRange regions,
                         SmallVectorImpl<Type> &inferredReturnTypes) {
   inferredReturnTypes.push_back(ShapeType::get(context));
   return success();
 }
 
 //===----------------------------------------------------------------------===//
-// AssumingOp
+// Assuming
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseAssumingOp(OpAsmParser &parser,
                                    OperationState &result) {
   result.regions.reserve(1);
   Region *doRegion = result.addRegion();
 
   auto &builder = parser.getBuilder();
if (yieldsResults) {
     p << " -> (" << op.getResultTypes() << ")";
   }
   p.printRegion(op.doRegion(),
                 /*printEntryBlockArgs=*/false,
                 /*printBlockTerminators=*/yieldsResults);
   p.printOptionalAttrDict(op.getAttrs());
 }
 
+namespace {
+// Removes AssumingOp with a passing condition and inlines the region.
+struct AssumingWithTrue : public OpRewritePattern<AssumingOp> {
+  using OpRewritePattern<AssumingOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AssumingOp op,
+                                PatternRewriter &rewriter) const override {
+    if (op.condition().getDefiningOp<TrueWitnessOp>()) {
+      auto *firstBlock = rewriter.getInsertionBlock();
+      auto *secondBlock = op.getBody();
+      auto initPosition = rewriter.getInsertionPoint();
+      auto *thirdBlock = rewriter.splitBlock(firstBlock, initPosition);
+
+      // Remove the AssumingOp and AssumingYieldOp.
+      auto &yieldOp = secondBlock->back();
+      rewriter.inlineRegionBefore(op.doRegion(), thirdBlock);
+      rewriter.replaceOp(op, yieldOp.getOperands());
+      rewriter.eraseOp(&yieldOp);
+
+      // Merge blocks together as there was no branching behavior from the
+      // AssumingOp.
+      rewriter.mergeBlocks(secondBlock, firstBlock);
+      rewriter.mergeBlocks(thirdBlock, firstBlock);
+      return success();
+    }
+    return failure();
+  }
+};
+}; // namespace
+
+void AssumingOp::getCanonicalizationPatterns(OwningRewritePatternList &patterns,
+                                             MLIRContext *context) {
+  // If taking true witness, inline region
+  patterns.insert<AssumingWithTrue>(context);
+}
+
+//===----------------------------------------------------------------------===//
+// 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,
ConstShapeOp::inferReturnTypes(MLIRContext *context,
                                Optional<Location> location, ValueRange operands,
                                DictionaryAttr attributes, RegionRange regions,
                                SmallVectorImpl<Type> &inferredReturnTypes) {
   inferredReturnTypes.push_back(ShapeType::get(context));
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// CstrBroadcastableOp
+//===----------------------------------------------------------------------===//
+
+namespace {
+// TODO(tpopp): Make this and other witness related operations commutative.
+// TODO(tpopp): Add a case for unknown shapes that are still defined by the same
+// operation.
+// TODO(tpopp): Once Witnesses are Attributes, replace this with folding.
+struct CstrBroadcastableTrue : public OpRewritePattern<CstrBroadcastableOp> {
+  using OpRewritePattern<CstrBroadcastableOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(CstrBroadcastableOp op,
+                                PatternRewriter &rewriter) const override {
+    // Don't try to compare equality when the shapes are not constant.
+    auto lhsShape = op.getOperand(0).getDefiningOp<ConstShapeOp>().shape();
+    auto rhsShape = op.getOperand(1).getDefiningOp<ConstShapeOp>().shape();
+    if (!(lhsShape && rhsShape))
+      return failure();
+
+    SmallVector<int64_t, 6> resultShape;
+    // If the shapes are not compatible, we can't fold it.
+    if (!OpTrait::util::getBroadcastedShape(
+            llvm::to_vector<6>(lhsShape.getValues<int64_t>()),
+            llvm::to_vector<6>(rhsShape.getValues<int64_t>()), resultShape))
+      return failure();
+
+    rewriter.replaceOpWithNewOp<TrueWitnessOp>(op);
+    return success();
+  }
+};
+}; // namespace
+
+void CstrBroadcastableOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
+  // If equal, return true op
+  patterns.insert<CstrBroadcastableTrue>(context);
+}
+
+//===----------------------------------------------------------------------===//
+// CstrEqOp
+//===----------------------------------------------------------------------===//
+
+void CstrEqOp::getCanonicalizationPatterns(OwningRewritePatternList &patterns,
+                                           MLIRContext *context) {
+  // If equal, return true op
+  patterns.insert<ConstCstrEq>(context);
+}
+
+//===----------------------------------------------------------------------===//
 // ConstSizeOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 ConstSizeOp::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>();
+        })
+      }]>>;
+
+def AllInputShapesEq : Constraint<CPred< [{
+        llvm::all_of($0.getOwner()->getOperands(), [&](mlir::Value val) {
+          return val.getDefiningOp<mlir::shape::ConstShapeOp>().shape() ==
+          $0.getOwner()->getOperand(0).getDefiningOp<mlir::shape::ConstShapeOp>().shape();
+        })
+      }]>>;
+
+// Canonicalization patterns.
+def ConstantAssumingAll : Pat<(Shape_AssumingAllOp:$op $input),
+          (Shape_TrueWitnessOp),
+          [(AllInputsTrueWitnesses $op)] >;
+
+def ConstCstrEq : Pat<(Shape_CstrEqOp:$op $shapes),
+          (Shape_TrueWitnessOp),
+          [(AllInputShapesEq $op)] >;
+
+

mlir/test/Dialect/Shape/canonicalize.mlir

 // Basic case.
 // CHECK-LABEL: func @f
 func @f() -> tensor<2xindex> {
   // CHECK: constant dense<[0, 1]> : tensor<2xindex>
   %cs = shape.const_shape [0, 1]
   %0 = "shape.to_extent_tensor"(%cs) : (!shape.shape) -> tensor<2xindex>
   return %0 : tensor<2xindex>
 }
+
+// -----
+// Confirm that cstr_eq, assuming_all, assuming, and any can all be folded away
+// when dependent only on constant shapes.
+// CHECK-LABEL: func @f
+func @f() -> !shape.shape {
+  // CHECK-NEXT: shape.const_shape
+  // CHECK-NEXT: return
+  %cs0 = shape.const_shape [0, 1]
+  %cs1 = shape.const_shape [0, 1]
+  %cs2 = shape.const_shape [0, 1]
+  %0 = shape.cstr_eq(%cs0, %cs1)
+  %1 = shape.cstr_eq(%cs0, %cs1, %cs2)
+  %2 = shape.assuming_all(%0, %1)
+  %3 = "shape.assuming"(%2) ({
+    %cs = shape.any(%cs0, %cs1, %cs2)
+    shape.assuming_yield %cs : !shape.shape
+  }) : (!shape.witness) -> !shape.shape
+  return %3 : !shape.shape
+}
+
+// -----
+// Confirm that cstr_broadcastable of const shapes is folded away.
+// CHECK-LABEL: func @f
+func @f() -> !shape.shape {
+  // CHECK: %[[CS:.*]] = shape.const_shape [1, 5]
+  // CHECK-NEXT: return %[[CS]]
+  %cs0 = shape.const_shape [1, 1]
+  %cs1 = shape.const_shape [1, 5]
+  %0 = shape.cstr_broadcastable(%cs0, %cs1)
+  %3 = "shape.assuming"(%0) ({
+    shape.assuming_yield %cs1 : !shape.shape
+  }) : (!shape.witness) -> !shape.shape
+  return %3 : !shape.shape
+}

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
 }