[MLIR] Canonicalize `shape.assuming` op to yield only inner values

[frgossen]

Summary

Depends On D99153

[bkramer]

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

273

Mind adding a comment on what this pattern does?

286

Won't this break if there are nested blocks within the assume?

[frgossen]

Address comments

[frgossen]

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

286

Thanks!

[bkramer]

Looks good now. I can't comment on whether there should be a more general pass to do this kind of transform (IfOp has the same issue), though.

[herhut]

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

273

Here is a counter example:

%w = cstr_eq(%shp0, %shp1)
%res = shape.assuming %w {
  %bcast = shape.broadcast(%shp0, %shp1)
  shape.yield %bcast
}

With the knowledge of %w being true, we can simplify the broadcast to %shp0 for example. Applying this canonicalization then would erase the fact that this is conditional on %w. As constraints are not side-effecting (on purpose), we would remove the constraint.

[frgossen]

Reverted https://github.com/llvm/llvm-project/commit/94ef248d7b76939cc3caacb83b7e168b82a74764

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

 //===- Shape.cpp - MLIR Shape Operations ----------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Shape/IR/Shape.h"
 
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Traits.h"
+#include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/raw_ostream.h"
LogicalResult matchAndRewrite(AssumingOp op,
     auto witness = op.witness().getDefiningOp<ConstWitnessOp>();
     if (!witness || !witness.passingAttr())
       return failure();
 
     AssumingOp::inlineRegionIntoParent(op, rewriter);
     return success();
   }
 };
+
+// Results of an assuming op that are defined outside its body are available
+// indepentently of the assuming op. There is no need to yield such values. This
+// canonicalization replaces such results with their definition.
+struct AssumingBypassIndependentResult : public OpRewritePattern<AssumingOp> {
+  using OpRewritePattern<AssumingOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AssumingOp op,
+                                PatternRewriter &rewriter) const override {
+    Block *body = op.getBody();
+    auto yieldOp = llvm::dyn_cast<AssumingYieldOp>(body->getTerminator());
+    if (!yieldOp)
+      return failure();
+
+    // See if there is at least one result that can bypass the assuming op.
+    auto isDefinedInBody = [&](Value val) {
+      Operation *def = val.getDefiningOp();
+      return def && op->isAncestor(def);
+    };
+    if (llvm::all_of(yieldOp.operands(), isDefinedInBody))
+      return failure();
+
+    SmallVector<Value, 2> replacementValues;
+    auto newAssumingOp = rewriter.create<shape::AssumingOp>(
+        op.getLoc(), op.witness(), [&](OpBuilder &b, Location loc) {
+          // Copy body.
+          BlockAndValueMapping mapping;
+          for (auto &nested : body->without_terminator())
+            b.clone(nested, mapping);
+
+          // Collect new yielded values.
+          SmallVector<Value, 2> mappedResults;
+          for (auto result : yieldOp.getOperands()) {
+            if (isDefinedInBody(result)) {
+              // This value is a result of the assuming op. We can obtain the
+              // replacement value only after the new op is fully constructed.
+              mappedResults.push_back(mapping.lookup(result));
+              replacementValues.push_back(nullptr);
+            } else {
+              // When defined outside of the assuming block, we can use it
+              // direclty. There is no need to yield the value from within the
+              // block.
+              replacementValues.push_back(result);
+            }
+          }
+          return mappedResults;
+        });
+
+    // Use the assuming op's results for the missing replacement values, which
+    // could not bypass the op.
+    auto src = newAssumingOp.getResults().begin();
+    for (auto &dst : replacementValues) {
+      if (dst)
+        continue;
+      dst = *src++;
+    }
+
+    rewriter.replaceOp(op, replacementValues);
+    return success();
+  }
+};
 } // namespace
 
 void AssumingOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                              MLIRContext *context) {
-  // If taking a passing witness, inline region.
-  patterns.add<AssumingWithTrue>(context);
+  patterns.add<AssumingBypassIndependentResult, AssumingWithTrue>(context);
 }
 
 // See RegionBranchOpInterface in Interfaces/ControlFlowInterfaces.td
 void AssumingOp::getSuccessorRegions(
     Optional<unsigned> index, ArrayRef<Attribute> operands,
     SmallVectorImpl<RegionSuccessor> &regions) {
   // AssumingOp has unconditional control flow into the region and back to the
   // parent, so return the correct RegionSuccessor purely based on the index

mlir/test/Dialect/Shape/canonicalize.mlir

 // CHECK-SAME: (%[[A:.*]]: tensor<3xindex>)
 func @broadcast_on_single_operand(%a : tensor<3xindex>) {
   // CHECK-NOT: broadcast
   // CHECK: "use"(%[[A]])
   %0 = shape.broadcast %a : tensor<3xindex> -> tensor<?xindex>
   "use"(%0) : (tensor<?xindex>) -> ()
   return
 }
+
+// -----
+
+// CHECK-LABEL: @bypass_assmunig
+// CHECK-SAME: (%[[ARG:.*]]: tensor<2x3xf32>)
+func @bypass_assmunig(%arg : tensor<2x3xf32>)
+    -> (tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>) {
+  // CHECK: %[[OUTER:.*]] = "some.tensor"
+  // CHECK: %[[WITNESS:.*]] = "some.witness"
+  // CHECK: %[[YIELDED:.*]] = shape.assuming %[[WITNESS]] -> (tensor<2x3xf32>) {
+  // CHECK:   %[[INNER:.*]] = "some.tensor"
+  // CHECK:   shape.assuming_yield %[[INNER]] : tensor<2x3xf32>
+  // CHECK: }
+  // CHECK: return %[[YIELDED]], %[[OUTER]], %[[ARG]]
+  %outer = "some.tensor"() : () -> tensor<2x3xf32>
+  %witness = "some.witness"() : () -> !shape.witness
+  %results:3 = shape.assuming %witness
+      -> (tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>) {
+    %inner = "some.tensor"() : () -> tensor<2x3xf32>
+    shape.assuming_yield %inner, %outer, %arg
+        : tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>
+  }
+  return %results#0, %results#1, %results#2
+      : tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>
+}