[mlir][SCF] Canonicalize dim(x) where x is an iter_arg

[springerm]

Summary

• Add DimOfIterArgFolder.
• Move existing cross-dialect canonicalization patterns to LoopCanonicalization.cpp.
• Rename SCFAffineOpCanonicalization pass to SCFForLoopCanonicalization.
• Expand documentaton of scf.for: The type of loop-carried variables may not change with iterations. (Not even the dynamic type.)

[ftynse]

mlir/lib/Dialect/SCF/Transforms/PassDetail.h

27

Please fix the linter warning.

[springerm]

address comments

mlir/include/mlir/Dialect/SCF/Passes.h

 std::unique_ptr<Pass> createForLoopSpecializationPass();
 
 /// Creates a pass that peels for loops at their upper bounds for
 /// better vectorization.
 std::unique_ptr<Pass> createForLoopPeelingPass();
 
 /// Creates a pass that canonicalizes affine.min and affine.max operations
 /// inside of scf.for loops with known lower and upper bounds.
-std::unique_ptr<Pass> createSCFAffineOpCanonicalizationPass();
+std::unique_ptr<Pass> createSCFForLoopCanonicalizationPass();
 
 /// Creates a loop fusion pass which fuses parallel loops.
 std::unique_ptr<Pass> createParallelLoopFusionPass();
 
 /// Creates a pass that specializes parallel loop for unrolling and
 /// vectorization.
 std::unique_ptr<Pass> createParallelLoopSpecializationPass();
 

mlir/include/mlir/Dialect/SCF/Passes.td

 include "mlir/Pass/PassBase.td"
 
 def SCFBufferize : FunctionPass<"scf-bufferize"> {
   let summary = "Bufferize the scf dialect.";
   let constructor = "mlir::createSCFBufferizePass()";
   let dependentDialects = ["memref::MemRefDialect"];
 }
 
-// Note: Making this a canonicalization pattern would require a dependency
-// of the SCF dialect on the Affine dialect or vice versa.
-def SCFAffineOpCanonicalization
-    : FunctionPass<"canonicalize-scf-affine-op"> {
-  let summary = "Canonicalize affine.min and affine.max ops in the context of "
-                "SCF loops with known bounds";
-  let constructor = "mlir::createSCFAffineOpCanonicalizationPass()";
-  let dependentDialects = ["AffineDialect"];
+// Note: Making these canonicalization patterns would require a dependency
+// of the SCF dialect on the Affine/Tensor/MemRef dialects or vice versa.
+def SCFForLoopCanonicalization
+    : FunctionPass<"for-loop-canonicalization"> {
+  let summary = "Canonicalize operations within scf.for loop bodies";
+  let constructor = "mlir::createSCFForLoopCanonicalizationPass()";
+  let dependentDialects = ["AffineDialect", "tensor::TensorDialect",
+                           "memref::MemRefDialect"];
 }
 
 def SCFForLoopPeeling
     : FunctionPass<"for-loop-peeling"> {
   let summary = "Peel `for` loops at their upper bounds.";
   let constructor = "mlir::createForLoopPeelingPass()";
   let options = [
     Option<"skipPartial", "skip-partial", "bool",

mlir/include/mlir/Dialect/SCF/SCFOps.td

 def ForOp : SCF_Op<"for",
       [DeclareOpInterfaceMethods<LoopLikeOpInterface>,
        DeclareOpInterfaceMethods<RegionBranchOpInterface>,
        SingleBlockImplicitTerminator<"scf::YieldOp">,
        RecursiveSideEffects]> {
   let summary = "for operation";
   let description = [{
     The "scf.for" operation represents a loop taking 3 SSA value as operands
-    that represent the lower bound, upper bound and step respectively.  The
+    that represent the lower bound, upper bound and step respectively. The
     operation defines an SSA value for its induction variable. It has one
     region capturing the loop body. The induction variable is represented as an
     argument of this region. This SSA value always has type index, which is the
     size of the machine word. The step is a value of type index, required to be
     positive.
     The lower and upper bounds specify a half-open range: the range includes
     the lower bound but does not include the upper bound.
 
     The body region must contain exactly one block that terminates with
     "scf.yield". Calling ForOp::build will create such a region and insert
     the terminator implicitly if none is defined, so will the parsing even in
     cases when it is absent from the custom format. For example:
 
     ```mlir
     scf.for %iv = %lb to %ub step %step {
       ... // body
     }
     ```
 
     `scf.for` can also operate on loop-carried variables and returns the final
     values after loop termination. The initial values of the variables are
     passed as additional SSA operands to the "scf.for" following the 3 loop
     control SSA values mentioned above (lower bound, upper bound and step). The
-    operation region has equivalent arguments for each variable representing
-    the value of the variable at the current iteration.
-
-    The region must terminate with a "scf.yield" that passes all the current
-    iteration variables to the next iteration, or to the "scf.for" result, if
-    at the last iteration. Note, that when the loop-carried variables are
-    present, calling ForOp::build will not insert the terminator implicitly.
-    The caller must insert "scf.yield" in that case.
+    operation region has an argument for the induction variable, followed by
+    one argument for each loop-carried variable, representing he value of the
+    variable at the current iteration.
+
+    The region must terminate with a "scf.yield" that passes the current
+    values of loop-carried variables to the next iteration, or to the "scf.for"
+    result, if at the last iteration. The type (static or dynamic) of a
+    loop-carried variable may not change with iterations. E.g., it is illegal
+    to pass a tensor of larger size to the next iteration; even if the tensor's
+    dimensions are dynamic (i.e., same static type). Note, that when the
+    loop-carried variables are present, calling ForOp::build will not insert the
+    terminator implicitly. The caller must insert "scf.yield" in that case.
 
     "scf.for" results hold the final values after the last iteration.
     For example, to sum-reduce a memref:
 
     ```mlir
     func @reduce(%buffer: memref<1024xf32>, %lb: index,
                  %ub: index, %step: index) -> (f32) {
       // Initial sum set to 0.

mlir/include/mlir/Dialect/SCF/Transforms.h

 /// S1(N) S2(N-1)                // Epilogue
 /// S2(N)                        // Epilogue
 void populateSCFLoopPipeliningPatterns(RewritePatternSet &patterns,
                                        const PipeliningOption &options);
 
 /// Populate patterns for canonicalizing operations inside SCF loop bodies.
 /// At the moment, only affine.min/max computations with iteration variables,
 /// loop bounds and loop steps are canonicalized.
-void populateSCFLoopBodyCanonicalizationPatterns(RewritePatternSet &patterns);
+void populateSCFForLoopCanonicalizationPatterns(RewritePatternSet &patterns);
 
 } // namespace scf
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SCF_TRANSFORMS_H_

mlir/lib/Dialect/Linalg/Transforms/CodegenStrategy.cpp

auto marker = (currentState == zeroState)
                       : linalg::LinalgTransformationFilter(
                             t->filter, currentState, nextState);
     stage1Patterns.emplace_back(t->buildRewritePatterns(context, marker));
     currentState = nextState;
   }
 
   RewritePatternSet stage2Patterns =
       linalg::getLinalgTilingCanonicalizationPatterns(context);
-  scf::populateSCFLoopBodyCanonicalizationPatterns(stage2Patterns);
+  scf::populateSCFForLoopCanonicalizationPatterns(stage2Patterns);
 
   auto stage3Transforms = [&](Operation *op) {
     // Some of these may be too aggressive as a stage 3 that is applied on each
     // stage 1 application and may have to be split out to post staged patterns
     // application (in which case they could just be passes, TBD).
     if (lateCodegenStrategyOptions.enableLICM) {
       op->walk([&](LoopLikeOpInterface loopLike) {
         LLVM_DEBUG(loopLike.print(llvm::dbgs() << "\nOriginal loop:\n"));

mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp

applyTilingToLoopPatterns(LinalgTilingLoopType loopType, FuncOp funcOp,
                           ArrayRef<StringRef> distributionTypes = {}) {
   auto options = LinalgTilingOptions()
                      .setTileSizes(tileSizes)
                      .setLoopType(loopType)
                      .setDistributionTypes(distributionTypes);
   MLIRContext *ctx = funcOp.getContext();
   RewritePatternSet patterns(ctx);
   insertTilingPatterns(patterns, options);
-  scf::populateSCFLoopBodyCanonicalizationPatterns(patterns);
+  scf::populateSCFForLoopCanonicalizationPatterns(patterns);
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
   (void)applyPatternsAndFoldGreedily(
       funcOp, getLinalgTilingCanonicalizationPatterns(ctx));
   // Drop the marker.
   funcOp.walk([](LinalgOp op) {
     op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
   });
 

mlir/lib/Dialect/SCF/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRSCFTransforms
   Bufferize.cpp
+  LoopCanonicalization.cpp
   LoopPipelining.cpp
   LoopRangeFolding.cpp
   LoopSpecialization.cpp
   ParallelLoopFusion.cpp
   ParallelLoopTiling.cpp
   StructuralTypeConversions.cpp
   Utils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SCF
 
   DEPENDS
   MLIRSCFPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRIR
   MLIRMemRef
   MLIRPass
   MLIRSCF
   MLIRStandard
   MLIRSupport
+  MLIRTensor
   MLIRTransforms
   MLIRTransformUtils
 )

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

+//===- LoopCanonicalization.cpp - Cross-dialect canonicalization patterns -===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains cross-dialect canonicalization patterns that cannot be
+// actual canonicalization patterns due to undesired additional dependencies.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/Passes.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/SCF/Transforms.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.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>
+///   ...
+/// }
+/// ```
+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();
+
+    Value initArg = forOp.getOpOperandForRegionIterArg(blockArg).get();
+    rewriter.updateRootInPlace(
+        dimOp, [&]() { dimOp.sourceMutable().assign(initArg); });
+
+    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 {
+    auto loopMatcher = [](Value iv, Value &lb, Value &ub, Value &step) {
+      if (scf::ForOp forOp = scf::getForInductionVarOwner(iv)) {
+        lb = forOp.lowerBound();
+        ub = forOp.upperBound();
+        step = forOp.step();
+        return success();
+      }
+      if (scf::ParallelOp parOp = scf::getParallelForInductionVarOwner(iv)) {
+        for (unsigned idx = 0; idx < parOp.getNumLoops(); ++idx) {
+          if (parOp.getInductionVars()[idx] == iv) {
+            lb = parOp.lowerBound()[idx];
+            ub = parOp.upperBound()[idx];
+            step = parOp.step()[idx];
+            return success();
+          }
+        }
+        return failure();
+      }
+      return failure();
+    };
+
+    return scf::canonicalizeMinMaxOpInLoop(rewriter, op, op.getAffineMap(),
+                                           op.operands(), IsMin, loopMatcher);
+  }
+};
+
+struct SCFForLoopCanonicalization
+    : public SCFForLoopCanonicalizationBase<SCFForLoopCanonicalization> {
+  void runOnFunction() override {
+    FuncOp funcOp = getFunction();
+    MLIRContext *ctx = funcOp.getContext();
+    RewritePatternSet patterns(ctx);
+    scf::populateSCFForLoopCanonicalizationPatterns(patterns);
+    if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(patterns))))
+      signalPassFailure();
+  }
+};
+} // namespace
+
+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);
+}
+
+std::unique_ptr<Pass> mlir::createSCFForLoopCanonicalizationPass() {
+  return std::make_unique<SCFForLoopCanonicalization>();
+}

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

struct ForLoopPeelingPattern : public OpRewritePattern<ForOp> {
 
   /// If set to true, loops inside partial iterations of another peeled loop
   /// are not peeled. This reduces the size of the generated code. Partial
   /// iterations are not usually performance critical.
   /// Note: Takes into account the entire chain of parent operations, not just
   /// the direct parent.
   bool skipPartial;
 };
-
-/// 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 {
-    auto loopMatcher = [](Value iv, Value &lb, Value &ub, Value &step) {
-      if (scf::ForOp forOp = scf::getForInductionVarOwner(iv)) {
-        lb = forOp.lowerBound();
-        ub = forOp.upperBound();
-        step = forOp.step();
-        return success();
-      }
-      if (scf::ParallelOp parOp = scf::getParallelForInductionVarOwner(iv)) {
-        for (unsigned idx = 0; idx < parOp.getNumLoops(); ++idx) {
-          if (parOp.getInductionVars()[idx] == iv) {
-            lb = parOp.lowerBound()[idx];
-            ub = parOp.upperBound()[idx];
-            step = parOp.step()[idx];
-            return success();
-          }
-        }
-        return failure();
-      }
-      return failure();
-    };
-
-    return scf::canonicalizeMinMaxOpInLoop(rewriter, op, op.getAffineMap(),
-                                           op.operands(), IsMin, loopMatcher);
-  }
-};
 } // namespace
 
 namespace {
 struct ParallelLoopSpecialization
     : public SCFParallelLoopSpecializationBase<ParallelLoopSpecialization> {
   void runOnFunction() override {
     getFunction().walk(
         [](ParallelOp op) { specializeParallelLoopForUnrolling(op); });
void runOnFunction() override {
 
     // Drop the markers.
     funcOp.walk([](Operation *op) {
       op->removeAttr(kPeeledLoopLabel);
       op->removeAttr(kPartialIterationLabel);
     });
   }
 };
-
-struct SCFAffineOpCanonicalization
-    : public SCFAffineOpCanonicalizationBase<SCFAffineOpCanonicalization> {
-  void runOnFunction() override {
-    FuncOp funcOp = getFunction();
-    MLIRContext *ctx = funcOp.getContext();
-    RewritePatternSet patterns(ctx);
-    scf::populateSCFLoopBodyCanonicalizationPatterns(patterns);
-    if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(patterns))))
-      signalPassFailure();
-  }
-};
 } // namespace
 
-std::unique_ptr<Pass> mlir::createSCFAffineOpCanonicalizationPass() {
-  return std::make_unique<SCFAffineOpCanonicalization>();
-}
-
 std::unique_ptr<Pass> mlir::createParallelLoopSpecializationPass() {
   return std::make_unique<ParallelLoopSpecialization>();
 }
 
 std::unique_ptr<Pass> mlir::createForLoopSpecializationPass() {
   return std::make_unique<ForLoopSpecialization>();
 }
 
 std::unique_ptr<Pass> mlir::createForLoopPeelingPass() {
   return std::make_unique<ForLoopPeeling>();
 }
-
-void mlir::scf::populateSCFLoopBodyCanonicalizationPatterns(
-    RewritePatternSet &patterns) {
-  MLIRContext *ctx = patterns.getContext();
-  patterns
-      .insert<AffineOpSCFCanonicalizationPattern<AffineMinOp, /*IsMin=*/true>,
-              AffineOpSCFCanonicalizationPattern<AffineMaxOp, /*IsMin=*/false>>(
-          ctx);
-}

mlir/lib/Dialect/SCF/Transforms/PassDetail.h

 void registerDialect(DialectRegistry &registry);
 
 class AffineDialect;
 
 namespace memref {
 class MemRefDialect;
 } // end namespace memref
 
+namespace tensor {
+class TensorDialect;
+}
+
 #define GEN_PASS_CLASSES
 #include "mlir/Dialect/SCF/Passes.h.inc"
 
 } // end namespace mlir
 
 #endif // DIALECT_LOOPOPS_TRANSFORMS_PASSDETAIL_H_

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

-// RUN: mlir-opt %s -canonicalize-scf-affine-op -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -for-loop-canonicalization -split-input-file | FileCheck %s
 
 // CHECK-LABEL: func @scf_for_canonicalize_min
 //       CHECK:   %[[C2:.*]] = constant 2 : i64
 //       CHECK:   scf.for
 //       CHECK:     memref.store %[[C2]], %{{.*}}[] : memref<i64>
 func @scf_for_canonicalize_min(%A : memref<i64>) {
   %c0 = constant 0 : index
   %c2 = constant 2 : index
func @scf_parallel_canonicalize_min_2(%A : memref<i64>) {
 
   scf.parallel (%i) = (%c1) to (%c7) step (%c2) {
     %1 = affine.min affine_map<(d0)[s0] -> (2, s0 - d0)> (%i)[%c7]
     %2 = index_cast %1: index to i64
     memref.store %2, %A[]: memref<i64>
   }
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @tensor_dim_of_iter_arg(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?x?xf32>
+//       CHECK:   scf.for
+//       CHECK:     tensor.dim %[[t]]
+func @tensor_dim_of_iter_arg(%t : 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 = %c0)
+      -> (tensor<?x?xf32>, index) {
+    %dim = tensor.dim %arg0, %c0 : tensor<?x?xf32>
+    scf.yield %arg0, %dim : tensor<?x?xf32>, index
+  }
+  return %1 : index
+}

mlir/test/lib/Dialect/Linalg/TestConvVectorization.cpp

void TestConvVectorization::runOnOperation() {
 
   SmallVector<RewritePatternSet, 4> stage1Patterns;
   linalg::populateConvVectorizationPatterns(context, stage1Patterns, tileSizes);
   SmallVector<FrozenRewritePatternSet, 4> frozenStage1Patterns;
   llvm::move(stage1Patterns, std::back_inserter(frozenStage1Patterns));
 
   RewritePatternSet stage2Patterns =
       linalg::getLinalgTilingCanonicalizationPatterns(context);
-  scf::populateSCFLoopBodyCanonicalizationPatterns(stage2Patterns);
+  scf::populateSCFForLoopCanonicalizationPatterns(stage2Patterns);
 
   auto stage3Transforms = [](Operation *op) {
     PassManager pm(op->getContext());
     pm.addPass(createLoopInvariantCodeMotionPass());
     if (failed(pm.run(cast<ModuleOp>(op))))
       llvm_unreachable("Unexpected failure in cleanup pass pipeline.");
     op->walk([](FuncOp func) {
       promoteSingleIterationLoops(func);

mlir/test/lib/Dialect/Linalg/TestLinalgFusionTransforms.cpp

struct TestLinalgGreedyFusion
   StringRef getDescription() const final {
     return "Test Linalg fusion by applying a greedy test transformation.";
   }
   void runOnFunction() override {
     MLIRContext *context = &getContext();
     RewritePatternSet patterns =
         linalg::getLinalgTilingCanonicalizationPatterns(context);
     patterns.add<ExtractSliceOfPadTensorSwapPattern>(context);
-    scf::populateSCFLoopBodyCanonicalizationPatterns(patterns);
+    scf::populateSCFForLoopCanonicalizationPatterns(patterns);
     FrozenRewritePatternSet frozenPatterns(std::move(patterns));
     do {
       (void)applyPatternsAndFoldGreedily(getFunction(), frozenPatterns);
       PassManager pm(context);
       pm.addPass(createLoopInvariantCodeMotionPass());
       pm.addPass(createCanonicalizerPass());
       pm.addPass(createCSEPass());
       LogicalResult res = pm.run(getFunction()->getParentOfType<ModuleOp>());

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":DialectUtils",
         ":IR",
         ":MemRefDialect",
         ":Pass",
         ":SCFDialect",
         ":SCFPassIncGen",
         ":StandardOps",
         ":Support",
+        ":TensorDialect",
         ":Transforms",
         "//llvm:Support",
     ],
 )
 
 ##---------------------------------------------------------------------------##
 # SparseTensor dialect.
 ##---------------------------------------------------------------------------##