[mlir][scf] Fold dim(scf.for) to dim(iter_arg)

[springerm]

Summary

Fold dim ops of scf.for results to dim ops of the respective iter args if the loop is shape preserving.

[herhut]

This is adding a specialized pattern that implements the shape reification interface + rewrite based on it. My first hunch was that this should then also be in a pass and not a canonicalization. However, I think control flow is different than compute ops in that regard, so I came to the conclusion that it makes sense to have this as a pattern.

The other comment would be that an interface might be helpful to model the shape preserving nature of operations and not hard-wire it to the current two ops. This CL did not introduce that part, so it should not be blocked on it.

Just noting this here to document that thought.

[springerm]

In D109430#2989101, @herhut wrote:

This is adding a specialized pattern that implements the shape reification interface + rewrite based on it. My first hunch was that this should then also be in a pass and not a canonicalization. However, I think control flow is different than compute ops in that regard, so I came to the conclusion that it makes sense to have this as a pattern.

Note: This is not actually a canonicalization pattern, but part of for-loop-canonicalization, which is a pass that contains "cross-dialect boundary" canonicalizations.

The other comment would be that an interface might be helpful to model the shape preserving nature of operations and not hard-wire it to the current two ops. This CL did not introduce that part, so it should not be blocked on it.

That's a good idea. Then this could be a canonicalization of tensor.dim/memref.dim. Same for D109431, which I just sent out.

Still not entirely sure how the interface would work, but maybe it could provide a function like:

/// Return the operand index, who's (dynamic) type is guaranteed to match with the `resultIdx`-th result's type. Return None if no such operand exists. 
Optional<unsigned> getTiedInput(unsigned resultIdx);

Or:

bool haveSameTypes(unsigned operandIdx, resultIdx);

Then it would be up to ops like scf.for to implement this interface. Ops such as tensor.insert_slice would also implement this interface, so the implementation of scf.for could query those without having a dependency on the tensor dialect.

[herhut]

In D109430#2989166, @springerm wrote:

bool haveSameTypes(unsigned operandIdx, resultIdx);

I like that variant, because the tiedInput concept is very linalg specific. I would just say that they have the same shape to cover the cases where element types are different. And SameOperandAndResultTypeand friends could provide a default implementation for this, too.

@mravishankar wdyt? I think this is far away enough from the other shape interfaces we have.

[springerm]

bool haveSameTypes(unsigned operandIdx, resultIdx);

I like that variant, because the tiedInput concept is very linalg specific. I would just say that they have the same shape to cover the cases where element types are different. And SameOperandAndResultTypeand friends could provide a default implementation for this, too.

If it's just the naming... One benefit of the getTiedInput variant (though it may have a different name) would be an easier implementation of isShapePreserving for scf.for loops. With haveSameTypes, I would have to iterate over all operands of an op, to see which operand has the same type. (There may even be multiple, so maybe the return value should be a SmallVector<unsigned>.)

Anyway, I'm landing this revision as is for the moment. I'll start a thread on Discourse or one of our chat groups later to discuss this further. And i'll change this implementation later if people are OK with adding a new interface.

mlir/lib/Dialect/SCF/Transforms/LoopCanonicalization.cpp

 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace mlir::scf;
 
-namespace {
-/// Fold dim ops of iter_args to dim ops of their respective init args. E.g.:
-///
-/// ```
-/// %0 = ... : tensor<?x?xf32>
-/// scf.for ... iter_args(%arg0 = %0) -> (tensor<?x?xf32>) {
-///   %1 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
-///   ...
-/// }
-/// ```
-///
-/// is folded to:
-///
-/// ```
-/// %0 = ... : tensor<?x?xf32>
-/// scf.for ... iter_args(%arg0 = %0) -> (tensor<?x?xf32>) {
-///   %1 = tensor.dim %0, %c0 : tensor<?x?xf32>
-///   ...
-/// }
-/// ```
-///
-/// Note: Dim ops are folded only if it can be proven that the runtime type of
-/// the iter arg does not change with loop iterations.
-template <typename OpTy>
-struct DimOfIterArgFolder : public OpRewritePattern<OpTy> {
-  using OpRewritePattern<OpTy>::OpRewritePattern;
-
-  /// A simple, conservative analysis to determine if the loop is shape
-  /// conserving. I.e., the type of the arg-th yielded value is the same as the
-  /// type of the corresponding basic block argument of the loop.
-  /// Note: This function handles only simple cases. Expand as needed.
-  static bool isShapePreserving(ForOp forOp, int64_t arg) {
-    auto yieldOp = cast<YieldOp>(forOp.getBody()->getTerminator());
-    assert(arg < static_cast<int64_t>(yieldOp.results().size()) &&
-           "arg is out of bounds");
-    Value value = yieldOp.results()[arg];
-    while (value) {
-      if (value == forOp.getRegionIterArgs()[arg])
-        return true;
-      OpResult opResult = value.dyn_cast<OpResult>();
-      if (!opResult)
-        return false;
 
-      using tensor::InsertSliceOp;
-      value =
-          llvm::TypeSwitch<Operation *, Value>(opResult.getOwner())
-              .template Case<InsertSliceOp>(
-                  [&](InsertSliceOp op) { return op.dest(); })
-              .template Case<ForOp>([&](ForOp forOp) {
-                return isShapePreserving(forOp, opResult.getResultNumber())
-                           ? forOp.getIterOperands()[opResult.getResultNumber()]
-                           : Value();
-              })
-              .Default([&](auto op) { return Value(); });
-    }
-    return false;
-  }
 
+namespace {
+/// Fold dim ops of iter_args to dim ops of their respective init args. E.g.:
+///
+/// ```
+/// %0 = ... : tensor<?x?xf32>
+/// scf.for ... iter_args(%arg0 = %0) -> (tensor<?x?xf32>) {
+///   %1 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
+///   ...
+/// }
+/// ```
+///
+/// is folded to:
+///
+/// ```
+/// %0 = ... : tensor<?x?xf32>
+/// scf.for ... iter_args(%arg0 = %0) -> (tensor<?x?xf32>) {
+///   %1 = tensor.dim %0, %c0 : tensor<?x?xf32>
+///   ...
+/// }
+/// ```
+///
+/// Note: Dim ops are folded only if it can be proven that the runtime type of
+/// the iter arg does not change with loop iterations.
+template <typename OpTy>
+struct DimOfIterArgFolder : public OpRewritePattern<OpTy> {
+  using OpRewritePattern<OpTy>::OpRewritePattern;
+
   LogicalResult matchAndRewrite(OpTy dimOp,
                                 PatternRewriter &rewriter) const override {
     auto blockArg = dimOp.source().template dyn_cast<BlockArgument>();
     if (!blockArg)
       return failure();
     auto forOp = dyn_cast<ForOp>(blockArg.getParentBlock()->getParentOp());
     if (!forOp)
       return failure();
     if (!isShapePreserving(forOp, blockArg.getArgNumber() - 1))
       return failure();
 
     Value initArg = forOp.getOpOperandForRegionIterArg(blockArg).get();
     rewriter.updateRootInPlace(
         dimOp, [&]() { dimOp.sourceMutable().assign(initArg); });
 
     return success();
   };
 };
 
+/// Fold dim ops of loop results to dim ops of their respective init args. E.g.:
+///
+/// ```
+/// %0 = ... : tensor<?x?xf32>
+/// %r = scf.for ... iter_args(%arg0 = %0) -> (tensor<?x?xf32>) {
+///   ...
+/// }
+/// %1 = tensor.dim %r, %c0 : tensor<?x?xf32>
+/// ```
+///
+/// is folded to:
+///
+/// ```
+/// %0 = ... : tensor<?x?xf32>
+/// %r = scf.for ... iter_args(%arg0 = %0) -> (tensor<?x?xf32>) {
+///   ...
+/// }
+/// %1 = tensor.dim %0, %c0 : tensor<?x?xf32>
+/// ```
+///
+/// Note: Dim ops are folded only if it can be proven that the runtime type of
+/// the iter arg does not change with loop iterations.
+template <typename OpTy>
+struct DimOfLoopResultFolder : public OpRewritePattern<OpTy> {
+  using OpRewritePattern<OpTy>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(OpTy dimOp,
+                                PatternRewriter &rewriter) const override {
+    auto forOp = dimOp.source().template getDefiningOp<scf::ForOp>();
+    if (!forOp)
+      return failure();
+    auto opResult = dimOp.source().template cast<OpResult>();
+    unsigned resultNumber = opResult.getResultNumber();
+    if (!isShapePreserving(forOp, resultNumber))
+      return failure();
+    rewriter.updateRootInPlace(dimOp, [&](){
+      dimOp.sourceMutable().assign(forOp.getIterOperands()[resultNumber]);
+    });
+    return success();
+  }
+};
+
 /// Canonicalize AffineMinOp/AffineMaxOp operations in the context of scf.for
 /// and scf.parallel loops with a known range.
 template <typename OpTy, bool IsMin>
 struct AffineOpSCFCanonicalizationPattern : public OpRewritePattern<OpTy> {
   using OpRewritePattern<OpTy>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(OpTy op,
                                 PatternRewriter &rewriter) const override {

 
 void mlir::scf::populateSCFForLoopCanonicalizationPatterns(
     RewritePatternSet &patterns) {
   MLIRContext *ctx = patterns.getContext();
   patterns
       .insert<AffineOpSCFCanonicalizationPattern<AffineMinOp, /*IsMin=*/true>,
               AffineOpSCFCanonicalizationPattern<AffineMaxOp, /*IsMin=*/false>,
               DimOfIterArgFolder<tensor::DimOp>,
-              DimOfIterArgFolder<memref::DimOp>>(ctx);
+              DimOfIterArgFolder<memref::DimOp>,
+              DimOfLoopResultFolder<tensor::DimOp>,
+              DimOfLoopResultFolder<memref::DimOp>>(ctx);
 }
 
 std::unique_ptr<Pass> mlir::createSCFForLoopCanonicalizationPass() {
   return std::make_unique<SCFForLoopCanonicalization>();
 }

mlir/test/Dialect/SCF/for-loop-canonicalization.mlir

func @tensor_dim_of_iter_arg_no_canonicalize(%t : tensor<?x?xf32>,
   %c10 = constant 10 : index
   %0, %1 = scf.for %i = %c0 to %c10 step %c1 iter_args(%arg0 = %t, %arg1 = %c0)
       -> (tensor<?x?xf32>, index) {
     %dim = tensor.dim %arg0, %c0 : tensor<?x?xf32>
     scf.yield %t2, %dim : tensor<?x?xf32>, index
   }
   return %1 : index
 }
+
+// -----
+
+// CHECK-LABEL: func @tensor_dim_of_loop_result(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?x?xf32>
+//       CHECK:   tensor.dim %[[t]]
+func @tensor_dim_of_loop_result(%t : tensor<?x?xf32>) -> index {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c10 = constant 10 : index
+  %0 = scf.for %i = %c0 to %c10 step %c1 iter_args(%arg0 = %t)
+      -> (tensor<?x?xf32>) {
+    scf.yield %arg0 : tensor<?x?xf32>
+  }
+  %dim = tensor.dim %0, %c0 : tensor<?x?xf32>
+  return %dim : index
+}
+
+// -----
+
+// CHECK-LABEL: func @tensor_dim_of_loop_result_no_canonicalize(
+//       CHECK:   %[[loop:.*]]:2 = scf.for
+//       CHECK:   tensor.dim %[[loop]]#1
+func @tensor_dim_of_loop_result_no_canonicalize(%t : tensor<?x?xf32>,
+                                                %u : tensor<?x?xf32>) -> index {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c10 = constant 10 : index
+  %0, %1 = scf.for %i = %c0 to %c10 step %c1 iter_args(%arg0 = %t, %arg1 = %u)
+      -> (tensor<?x?xf32>, tensor<?x?xf32>) {
+    scf.yield %arg0, %u : tensor<?x?xf32>, tensor<?x?xf32>
+  }
+  %dim = tensor.dim %1, %c0 : tensor<?x?xf32>
+  return %dim : index
+}