[mlir] Canonicalization of shape.assuming

[tpopp]

Summary

This will inline the region to a shape.assuming in the case that the
input witness is found to be statically true.

[jpienaar]

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

169

Early return?

171

Is there perhaps more descriptive names we could use here?

[herhut]

Thanks. I think this is ready to land with @jpienaar comments addressed.

[tpopp]

Rename variables.

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

let description = [{
     ready to execute.
   }];
   let arguments = (ins Shape_WitnessType:$witness);
   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/lib/Dialect/Shape/IR/Shape.cpp

if (yieldsResults) {
     p << " -> (" << op.getResultTypes() << ")";
   }
   p.printRegion(op.doRegion(),
                 /*printEntryBlockArgs=*/false,
                 /*printBlockTerminators=*/yieldsResults);
   p.printOptionalAttrDict(op.getAttrs());
 }
 
+namespace {
+// Removes AssumingOp with a passing witness and inlines the region.
+struct AssumingWithTrue : public OpRewritePattern<AssumingOp> {
+  using OpRewritePattern<AssumingOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AssumingOp op,
+                                PatternRewriter &rewriter) const override {
+    auto witness = op.witness().getDefiningOp<ConstWitnessOp>();
+    if (!witness || !witness.passingAttr())
+      return failure();
+
+    auto *blockBeforeAssuming = rewriter.getInsertionBlock();
+    auto *assumingBlock = op.getBody();
+    auto initPosition = rewriter.getInsertionPoint();
+    auto *blockAfterAssuming =
+        rewriter.splitBlock(blockBeforeAssuming, initPosition);
+
+    // Remove the AssumingOp and AssumingYieldOp.
+    auto &yieldOp = assumingBlock->back();
+    rewriter.inlineRegionBefore(op.doRegion(), blockAfterAssuming);
+    rewriter.replaceOp(op, yieldOp.getOperands());
+    rewriter.eraseOp(&yieldOp);
+
+    // Merge blocks together as there was no branching behavior from the
+    // AssumingOp.
+    rewriter.mergeBlocks(assumingBlock, blockBeforeAssuming);
+    rewriter.mergeBlocks(blockAfterAssuming, blockBeforeAssuming);
+    return success();
+  }
+};
+}; // namespace
+
+void AssumingOp::getCanonicalizationPatterns(OwningRewritePatternList &patterns,
+                                             MLIRContext *context) {
+  // If taking a passing witness, inline region
+  patterns.insert<AssumingWithTrue>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // 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];

mlir/test/Dialect/Shape/canonicalize.mlir

 func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> !shape.shape {
   // CHECK-NEXT: shape.any
   // CHECK-NEXT: return %[[CS]]
   %1 = shape.any %arg0, %arg1
   return %1 : !shape.shape
 }
 
 // -----
+// assuming with a known passing witness can be removed
+// CHECK-LABEL: func @f
+func @f() {
+  // CHECK-NEXT: source
+  // CHECK-NEXT: sink
+  // CHECK-NEXT: return
+  %0 = shape.const_witness true
+  %1 = shape.assuming %0 -> index {
+    %2 = "test.source"() : () -> (index)
+    shape.assuming_yield %2 : index
+  }
+  "test.sink"(%1) : (index) -> ()
+  return
+}
+
+// -----
+// assuming without a known passing passing witness cannot be removed
+// CHECK-LABEL: func @f
+func @f() {
+  // CHECK-NEXT: test.source
+  // CHECK-NEXT: shape.assuming
+  // CHECK-NEXT:   test.source
+  // CHECK-NEXT:   shape.assuming_yield
+  // CHECK-NEXT: }
+  // CHECK-NEXT: test.sink
+  // CHECK-NEXT: return
+  %0 = "test.source"() : () -> (!shape.witness)
+  %1 = shape.assuming %0 -> index {
+    %2 = "test.source"() : () -> (index)
+    shape.assuming_yield %2 : index
+  }
+  "test.sink"(%1) : (index) -> ()
+  return
+}
+
+// -----
 // Broadcastable with broadcastable constant shapes can be removed.
 // CHECK-LABEL: func @f
 func @f() {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %cs0 = shape.const_shape [3, 1]
   %cs1 = shape.const_shape [1, 5]