[Mlir-commits] [mlir] [mlir] Add loop bounds normalization pass (PR #93781)
Jorn Tuyls
llvmlistbot at llvm.org
Thu May 30 01:10:41 PDT 2024
https://github.com/jtuyls created https://github.com/llvm/llvm-project/pull/93781
Add pass to normalize loop bounds, i.e. i.e. calculate new normalized upper bounds for lower bounds equal to zero and step sizes equal to one. Then, insert new `affine.apply` operations to calculate the denormalized index values and update all usage from the original induction variables to the results of the `affine.apply` operations.
I created a new interface for loop-like operations with induction variables (`LoopLikeWithInductionVarsOpInterface`) instead of putting the new methods in `LoopLikeOpInterface` as not all loop-like operations have induction variables and therefore wouldn't want to implement this interface.
cc @MaheshRavishankar @qedawkins
>From 1fe405faf6c09d0ff0d86875584b9dd864479e48 Mon Sep 17 00:00:00 2001
From: Jorn Tuyls <jorn.tuyls at gmail.com>
Date: Thu, 30 May 2024 00:45:10 -0700
Subject: [PATCH] [mlir] Add loop bound normalization pass
---
mlir/include/mlir/Dialect/SCF/IR/SCFOps.td | 97 +++++++
mlir/include/mlir/Dialect/Utils/LoopUtils.h | 30 ++
.../mlir/Interfaces/LoopLikeInterface.h | 4 +
.../mlir/Interfaces/LoopLikeInterface.td | 126 +++++++++
mlir/include/mlir/Transforms/Passes.h | 4 +
mlir/include/mlir/Transforms/Passes.td | 6 +
mlir/lib/Dialect/SCF/IR/SCF.cpp | 60 ++++
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 55 +---
mlir/lib/Dialect/Utils/CMakeLists.txt | 1 +
mlir/lib/Dialect/Utils/LoopUtils.cpp | 52 ++++
mlir/lib/IR/Operation.cpp | 2 -
mlir/lib/Interfaces/LoopLikeInterface.cpp | 25 ++
mlir/lib/Transforms/CMakeLists.txt | 2 +
mlir/lib/Transforms/NormalizeLoopBounds.cpp | 118 ++++++++
mlir/test/Dialect/Affine/loop-coalescing.mlir | 11 +-
.../Transforms/normalize-loop-bounds.mlir | 266 ++++++++++++++++++
16 files changed, 796 insertions(+), 63 deletions(-)
create mode 100644 mlir/include/mlir/Dialect/Utils/LoopUtils.h
create mode 100644 mlir/lib/Dialect/Utils/LoopUtils.cpp
create mode 100644 mlir/lib/Transforms/NormalizeLoopBounds.cpp
create mode 100644 mlir/test/Transforms/normalize-loop-bounds.mlir
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 0b063aa772bab..0e23257456223 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -140,6 +140,7 @@ def ForOp : SCF_Op<"for",
"getSingleUpperBound", "getYieldedValuesMutable",
"promoteIfSingleIteration", "replaceWithAdditionalYields",
"yieldTiledValuesAndReplace"]>,
+ LoopLikeWithInductionVarsOpInterface,
AllTypesMatch<["lowerBound", "upperBound", "step"]>,
ConditionallySpeculatable,
DeclareOpInterfaceMethods<RegionBranchOpInterface,
@@ -267,6 +268,74 @@ def ForOp : SCF_Op<"for",
return getBody()->getArguments().drop_front(getNumInductionVars())[index];
}
+ /// Return the induction variables.
+ ::mlir::ValueRange getInductionVars() {
+ return getBody()->getArguments().take_front(getNumInductionVars());
+ }
+
+ /// Get lower bounds as `OpFoldResult`.
+ SmallVector<OpFoldResult> getMixedLowerBound() {
+ return {getAsOpFoldResult(getLowerBound())};
+ }
+
+ /// Get upper bounds as `OpFoldResult`.
+ SmallVector<OpFoldResult> getMixedUpperBound() {
+ return {getAsOpFoldResult(getUpperBound())};
+ }
+
+ // Get steps as `OpFoldResult`.
+ SmallVector<OpFoldResult> getMixedStep() {
+ return {getAsOpFoldResult(getStep())};
+ }
+
+ /// Get lower bounds as values.
+ SmallVector<Value> getLowerBound(OpBuilder &b) {
+ return ValueRange{getLowerBound()};
+ }
+
+ /// Get upper bounds as values.
+ SmallVector<Value> getUpperBound(OpBuilder &b) {
+ return ValueRange{getUpperBound()};
+ }
+
+ /// Get steps as values.
+ SmallVector<Value> getStep(OpBuilder &b) {
+ return ValueRange{getStep()};
+ }
+
+ /// Set the lower bounds from `OpFoldResult`.
+ void setMixedLowerBounds(OpBuilder &b, ArrayRef<OpFoldResult> lbs) {
+ setLowerBound(getValueOrCreateConstantIndexOp(b, getLoc(), lbs[0]));
+ }
+
+ /// Set the upper bounds from `OpFoldResult`.
+ void setMixedUpperBounds(OpBuilder &b, ArrayRef<OpFoldResult> ubs) {
+ setUpperBound(getValueOrCreateConstantIndexOp(b, getLoc(), ubs[0]));
+ }
+
+ /// Set the steps from `OpFoldResult`.
+ void setMixedSteps(OpBuilder &b, ArrayRef<OpFoldResult> steps) {
+ setStep(getValueOrCreateConstantIndexOp(b, getLoc(), steps[0]));
+ }
+
+ /// Set the lower bounds from values.
+ void setLowerBounds(ArrayRef<Value> lbs) {
+ assert(lbs.size() == 1 && "expected a single lower bound");
+ setLowerBound(lbs[0]);
+ }
+
+ /// Set the upper bounds from values.
+ void setUpperBounds(ArrayRef<Value> ubs) {
+ assert(ubs.size() == 1 && "expected a single upper bound");
+ setUpperBound(ubs[0]);
+ }
+
+ /// Set the steps from values.
+ void setSteps(ArrayRef<Value> steps) {
+ assert(steps.size() == 1 && "expected a single step");
+ setStep(steps[0]);
+ }
+
void setLowerBound(Value bound) { getOperation()->setOperand(0, bound); }
void setUpperBound(Value bound) { getOperation()->setOperand(1, bound); }
void setStep(Value step) { getOperation()->setOperand(2, step); }
@@ -304,6 +373,7 @@ def ForallOp : SCF_Op<"forall", [
["getInitsMutable", "getRegionIterArgs", "getSingleInductionVar",
"getSingleLowerBound", "getSingleUpperBound", "getSingleStep",
"promoteIfSingleIteration", "yieldTiledValuesAndReplace"]>,
+ LoopLikeWithInductionVarsOpInterface,
RecursiveMemoryEffects,
SingleBlockImplicitTerminator<"scf::InParallelOp">,
DeclareOpInterfaceMethods<RegionBranchOpInterface>,
@@ -543,6 +613,33 @@ def ForallOp : SCF_Op<"forall", [
return getValueOrCreateConstantIndexOp(b, getLoc(), getMixedStep());
}
+ /// Set the lower bounds from `OpFoldResult`.
+ void setMixedLowerBounds(OpBuilder &b, ArrayRef<OpFoldResult> lbs);
+
+ /// Set the upper bounds from `OpFoldResult`.
+ void setMixedUpperBounds(OpBuilder &b, ArrayRef<OpFoldResult> ubs);
+
+ /// Set the steps from `OpFoldResult`.
+ void setMixedSteps(OpBuilder &b, ArrayRef<OpFoldResult> steps);
+
+ /// Set the lower bounds from values.
+ void setLowerBounds(ArrayRef<Value> lbs) {
+ OpBuilder b(getOperation()->getContext());
+ return setMixedLowerBounds(b, getAsOpFoldResult(lbs));
+ }
+
+ /// Set the upper bounds from values.
+ void setUpperBounds(ArrayRef<Value> ubs) {
+ OpBuilder b(getOperation()->getContext());
+ return setMixedUpperBounds(b, getAsOpFoldResult(ubs));
+ }
+
+ /// Set the steps from values.
+ void setSteps(ArrayRef<Value> steps) {
+ OpBuilder b(getOperation()->getContext());
+ return setMixedSteps(b, getAsOpFoldResult(steps));
+ }
+
int64_t getRank() { return getStaticLowerBound().size(); }
/// Number of operands controlling the loop: lbs, ubs, steps
diff --git a/mlir/include/mlir/Dialect/Utils/LoopUtils.h b/mlir/include/mlir/Dialect/Utils/LoopUtils.h
new file mode 100644
index 0000000000000..15e901dc0e45e
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Utils/LoopUtils.h
@@ -0,0 +1,30 @@
+//===- LoopUtils.h - Helpers related to loop operations ---------*- C++ -*-===//
+//
+// 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 header file defines utilities for loop operations.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/PatternMatch.h"
+
+namespace mlir {
+
+// This structure is to pass and return sets of loop parameters without
+// confusing the order.
+struct LoopParams {
+ Value lowerBound;
+ Value upperBound;
+ Value step;
+};
+
+/// Calculate the normalized loop upper bounds with lower bound equal to zero
+/// and step equal to one.
+LoopParams emitNormalizedLoopBounds(RewriterBase &rewriter, Location loc,
+ Value lb, Value ub, Value step);
+
+} // namespace mlir
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.h b/mlir/include/mlir/Interfaces/LoopLikeInterface.h
index 42609e824c86a..fab5ffa26e574 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.h
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.h
@@ -13,6 +13,7 @@
#ifndef MLIR_INTERFACES_LOOPLIKEINTERFACE_H_
#define MLIR_INTERFACES_LOOPLIKEINTERFACE_H_
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/OpDefinition.h"
namespace mlir {
@@ -28,6 +29,9 @@ using NewYieldValuesFn = std::function<SmallVector<Value>(
namespace detail {
/// Verify invariants of the LoopLikeOpInterface.
LogicalResult verifyLoopLikeOpInterface(Operation *op);
+
+/// Verify invariants of the LoopLikeWithInductionVarsOpInterface.
+LogicalResult verifyLoopLikeWithInductionVarsOpInterface(Operation *op);
} // namespace detail
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
index f0dc6e60eba58..95a8c5a244b62 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -375,6 +375,132 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
}];
}
+def LoopLikeWithInductionVarsOpInterface
+ : OpInterface<"LoopLikeWithInductionVarsOpInterface"> {
+ let description = [{
+ Interface for loop-like operations with one or more induction variables.
+ This interface contains helper functions for retrieving and updating the
+ lower bound, upper bound and step size for each induction variable and
+ provides a utility function to check whether the loop is normalized., i.e.
+ all lower bounds are equal to zero and steps are equal to one.
+ }];
+ let cppNamespace = "::mlir";
+
+ let methods = [
+ InterfaceMethod<[{
+ Return the induction variables if they exist, otherwise return
+ std::nullopt.
+ }],
+ /*retTy=*/"::mlir::ValueRange",
+ /*methodName=*/"getInductionVars"
+ >,
+ InterfaceMethod<[{
+ Return the lower bound values or attributes as OpFoldResult.
+ }],
+ /*retTy=*/"SmallVector<::mlir::OpFoldResult>",
+ /*methodName=*/"getMixedLowerBound"
+ >,
+ InterfaceMethod<[{
+ Return the step values or attributes if they exist as OpFoldResult.
+ }],
+ /*retTy=*/"SmallVector<::mlir::OpFoldResult>",
+ /*methodName=*/"getMixedStep"
+ >,
+ InterfaceMethod<[{
+ Return the upper bound values or attributes as OpFoldResult.
+ }],
+ /*retTy=*/"SmallVector<::mlir::OpFoldResult>",
+ /*methodName=*/"getMixedUpperBound"
+ >,
+ InterfaceMethod<[{
+ Return the lower bounds as values.
+ }],
+ /*retTy=*/"SmallVector<Value>",
+ /*methodName=*/"getLowerBound",
+ /*args=*/(ins "OpBuilder &":$b)
+ >,
+ InterfaceMethod<[{
+ Return the steps as values.
+ }],
+ /*retTy=*/"SmallVector<Value>",
+ /*methodName=*/"getStep",
+ /*args=*/(ins "OpBuilder &":$b)
+ >,
+ InterfaceMethod<[{
+ Return the upper bounds as values.
+ }],
+ /*retTy=*/"SmallVector<Value>",
+ /*methodName=*/"getUpperBound",
+ /*args=*/(ins "OpBuilder &":$b)
+ >,
+ InterfaceMethod<[{
+ Set the lower bounds from an array of `OpFoldResult`.
+ }],
+ /*retTy=*/"void",
+ /*methodName=*/"setMixedLowerBounds",
+ /*args=*/(ins "OpBuilder &":$b, "ArrayRef<OpFoldResult>":$lbs)
+ >,
+ InterfaceMethod<[{
+ Set the steps from an array of `OpFoldResult`.
+ }],
+ /*retTy=*/"void",
+ /*methodName=*/"setMixedSteps",
+ /*args=*/(ins "OpBuilder &":$b, "ArrayRef<OpFoldResult>":$lbs)
+ >,
+ InterfaceMethod<[{
+ Set the upper bounds from an array of `OpFoldResult`.
+ }],
+ /*retTy=*/"void",
+ /*methodName=*/"setMixedUpperBounds",
+ /*args=*/(ins "OpBuilder &":$b, "ArrayRef<OpFoldResult>":$lbs)
+ >,
+ InterfaceMethod<[{
+ Set the lower bounds from an array of values.
+ }],
+ /*retTy=*/"void",
+ /*methodName=*/"setLowerBounds",
+ /*args=*/(ins "ArrayRef<Value>":$lbs)
+ >,
+ InterfaceMethod<[{
+ Set the steps from an array of values.
+ }],
+ /*retTy=*/"void",
+ /*methodName=*/"setSteps",
+ /*args=*/(ins "ArrayRef<Value>":$lbs)
+ >,
+ InterfaceMethod<[{
+ Set the upper bounds from an array of values.
+ }],
+ /*retTy=*/"void",
+ /*methodName=*/"setUpperBounds",
+ /*args=*/(ins "ArrayRef<Value>":$lbs)
+ >,
+ InterfaceMethod<[{
+ Checks if the lower bounds are zeros and steps are ones.
+ }],
+ /*retTy=*/"bool",
+ /*methodName=*/"isNormalized",
+ /*args=*/(ins),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ auto allEqual = [](ArrayRef<OpFoldResult> results, int64_t val) {
+ return llvm::all_of(results, [&](OpFoldResult ofr) {
+ auto intValue = getConstantIntValue(ofr);
+ return intValue.has_value() && intValue == val;
+ });
+ };
+ SmallVector<::mlir::OpFoldResult> lbs = $_op.getMixedLowerBound();
+ SmallVector<::mlir::OpFoldResult> steps = $_op.getMixedStep();
+ return allEqual(lbs, 0) && allEqual(steps, 1);
+ }]
+ >
+ ];
+
+ let verify = [{
+ return detail::verifyLoopLikeWithInductionVarsOpInterface($_op);
+ }];
+}
+
//===----------------------------------------------------------------------===//
// Traits
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Transforms/Passes.h b/mlir/include/mlir/Transforms/Passes.h
index 58bd61b2ae8b8..755ec7ecdfbad 100644
--- a/mlir/include/mlir/Transforms/Passes.h
+++ b/mlir/include/mlir/Transforms/Passes.h
@@ -82,6 +82,10 @@ std::unique_ptr<Pass> createLoopInvariantCodeMotionPass();
/// Creates a pass that hoists loop-invariant subset ops.
std::unique_ptr<Pass> createLoopInvariantSubsetHoistingPass();
+/// Create a pass that normalizes the loop bounds of loop-like operations with
+/// induction variables.
+std::unique_ptr<Pass> createNormalizeLoopBoundsPass();
+
/// Creates a pass to strip debug information from a function.
std::unique_ptr<Pass> createStripDebugInfoPass();
diff --git a/mlir/include/mlir/Transforms/Passes.td b/mlir/include/mlir/Transforms/Passes.td
index 1b40a87c63f27..5d1256e502a12 100644
--- a/mlir/include/mlir/Transforms/Passes.td
+++ b/mlir/include/mlir/Transforms/Passes.td
@@ -377,6 +377,12 @@ def Mem2Reg : Pass<"mem2reg"> {
];
}
+def NormalizeLoopBounds : Pass<"normalize-loop-bounds"> {
+ let summary = "Normalize the loop bounds of loop-like operations with "
+ "induction variables.";
+ let constructor = "mlir::createNormalizeLoopBoundsPass()";
+}
+
def PrintOpStats : Pass<"print-op-stats"> {
let summary = "Print statistics of operations";
let constructor = "mlir::createPrintOpStatsPass()";
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 107fd0690f193..3e7becb094b6b 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -1387,6 +1387,66 @@ void ForallOp::build(
build(b, result, lbs, ubs, steps, outputs, mapping, bodyBuilderFn);
}
+/// Set the lower bounds from `OpFoldResult`.
+void ForallOp::setMixedLowerBounds(OpBuilder &b, ArrayRef<OpFoldResult> lbs) {
+ SmallVector<int64_t> staticLbs;
+ SmallVector<Value> dynamicLbs;
+ dispatchIndexOpFoldResults(lbs, dynamicLbs, staticLbs);
+ getOperation()->setOperands(0, getDynamicLowerBound().size(), dynamicLbs);
+ (*this)->setAttr(getStaticLowerBoundAttrName(),
+ b.getDenseI64ArrayAttr(staticLbs));
+ ArrayRef<int32_t> segmentSizes =
+ (*this)
+ ->getAttrOfType<DenseI32ArrayAttr>("operandSegmentSizes")
+ .asArrayRef();
+ SmallVector<int32_t> newSegmentSizes(segmentSizes.begin(),
+ segmentSizes.end());
+ newSegmentSizes[0] = dynamicLbs.size();
+ (*this)->setAttr("operandSegmentSizes",
+ b.getDenseI32ArrayAttr(newSegmentSizes));
+}
+
+/// Set the upper bounds from `OpFoldResult`.
+void ForallOp::setMixedUpperBounds(OpBuilder &b, ArrayRef<OpFoldResult> ubs) {
+ SmallVector<int64_t> staticUbs;
+ SmallVector<Value> dynamicUbs;
+ dispatchIndexOpFoldResults(ubs, dynamicUbs, staticUbs);
+ size_t offset = getDynamicLowerBound().size();
+ getOperation()->setOperands(offset, getDynamicUpperBound().size(),
+ dynamicUbs);
+ (*this)->setAttr(getStaticUpperBoundAttrName(),
+ b.getDenseI64ArrayAttr(staticUbs));
+ ArrayRef<int32_t> segmentSizes =
+ (*this)
+ ->getAttrOfType<DenseI32ArrayAttr>("operandSegmentSizes")
+ .asArrayRef();
+ SmallVector<int32_t> newSegmentSizes(segmentSizes.begin(),
+ segmentSizes.end());
+ newSegmentSizes[1] = dynamicUbs.size();
+ (*this)->setAttr("operandSegmentSizes",
+ b.getDenseI32ArrayAttr(newSegmentSizes));
+}
+
+/// Set the steps from `OpFoldResult`.
+void ForallOp::setMixedSteps(OpBuilder &b, ArrayRef<OpFoldResult> steps) {
+ SmallVector<int64_t> staticSteps;
+ SmallVector<Value> dynamicSteps;
+ dispatchIndexOpFoldResults(steps, dynamicSteps, staticSteps);
+ size_t offset = getDynamicLowerBound().size() + getDynamicUpperBound().size();
+ getOperation()->setOperands(offset, getDynamicStep().size(), dynamicSteps);
+ (*this)->setAttr(getStaticStepAttrName(),
+ b.getDenseI64ArrayAttr(staticSteps));
+ ArrayRef<int32_t> segmentSizes =
+ (*this)
+ ->getAttrOfType<DenseI32ArrayAttr>("operandSegmentSizes")
+ .asArrayRef();
+ SmallVector<int32_t> newSegmentSizes(segmentSizes.begin(),
+ segmentSizes.end());
+ newSegmentSizes[2] = dynamicSteps.size();
+ (*this)->setAttr("operandSegmentSizes",
+ b.getDenseI32ArrayAttr(newSegmentSizes));
+}
+
// Checks if the lbs are zeros and steps are ones.
bool ForallOp::isNormalized() {
auto allEqual = [](ArrayRef<OpFoldResult> results, int64_t val) {
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 6658cca03eba7..41f52cb84f4ed 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -16,6 +16,7 @@
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Utils/LoopUtils.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/PatternMatch.h"
@@ -29,16 +30,6 @@
using namespace mlir;
-namespace {
-// This structure is to pass and return sets of loop parameters without
-// confusing the order.
-struct LoopParams {
- Value lowerBound;
- Value upperBound;
- Value step;
-};
-} // namespace
-
SmallVector<scf::ForOp> mlir::replaceLoopNestWithNewYields(
RewriterBase &rewriter, MutableArrayRef<scf::ForOp> loopNest,
ValueRange newIterOperands, const NewYieldValuesFn &newYieldValuesFn,
@@ -473,50 +464,6 @@ LogicalResult mlir::loopUnrollByFactor(
return success();
}
-/// Transform a loop with a strictly positive step
-/// for %i = %lb to %ub step %s
-/// into a 0-based loop with step 1
-/// for %ii = 0 to ceildiv(%ub - %lb, %s) step 1 {
-/// %i = %ii * %s + %lb
-/// Insert the induction variable remapping in the body of `inner`, which is
-/// expected to be either `loop` or another loop perfectly nested under `loop`.
-/// Insert the definition of new bounds immediate before `outer`, which is
-/// expected to be either `loop` or its parent in the loop nest.
-static LoopParams emitNormalizedLoopBounds(RewriterBase &rewriter, Location loc,
- Value lb, Value ub, Value step) {
- // For non-index types, generate `arith` instructions
- // Check if the loop is already known to have a constant zero lower bound or
- // a constant one step.
- bool isZeroBased = false;
- if (auto lbCst = getConstantIntValue(lb))
- isZeroBased = lbCst.value() == 0;
-
- bool isStepOne = false;
- if (auto stepCst = getConstantIntValue(step))
- isStepOne = stepCst.value() == 1;
-
- // Compute the number of iterations the loop executes: ceildiv(ub - lb, step)
- // assuming the step is strictly positive. Update the bounds and the step
- // of the loop to go from 0 to the number of iterations, if necessary.
- if (isZeroBased && isStepOne)
- return {lb, ub, step};
-
- Value diff = isZeroBased ? ub : rewriter.create<arith::SubIOp>(loc, ub, lb);
- Value newUpperBound =
- isStepOne ? diff : rewriter.create<arith::CeilDivSIOp>(loc, diff, step);
-
- Value newLowerBound = isZeroBased
- ? lb
- : rewriter.create<arith::ConstantOp>(
- loc, rewriter.getZeroAttr(lb.getType()));
- Value newStep = isStepOne
- ? step
- : rewriter.create<arith::ConstantOp>(
- loc, rewriter.getIntegerAttr(step.getType(), 1));
-
- return {newLowerBound, newUpperBound, newStep};
-}
-
/// Get back the original induction variable values after loop normalization
static void denormalizeInductionVariable(RewriterBase &rewriter, Location loc,
Value normalizedIv, Value origLb,
diff --git a/mlir/lib/Dialect/Utils/CMakeLists.txt b/mlir/lib/Dialect/Utils/CMakeLists.txt
index a0096e5f299d5..41b2fe287beb3 100644
--- a/mlir/lib/Dialect/Utils/CMakeLists.txt
+++ b/mlir/lib/Dialect/Utils/CMakeLists.txt
@@ -1,5 +1,6 @@
add_mlir_library(MLIRDialectUtils
IndexingUtils.cpp
+ LoopUtils.cpp
ReshapeOpsUtils.cpp
StructuredOpsUtils.cpp
StaticValueUtils.cpp
diff --git a/mlir/lib/Dialect/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Utils/LoopUtils.cpp
new file mode 100644
index 0000000000000..3d8aa5ef7dfc1
--- /dev/null
+++ b/mlir/lib/Dialect/Utils/LoopUtils.cpp
@@ -0,0 +1,52 @@
+//===- LoopUtils.cpp - Helpers related to loop 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/Utils/LoopUtils.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+
+using namespace mlir;
+
+/// Calculate the normalized loop upper bounds with lower bound equal to zero
+/// and step equal to one.
+LoopParams mlir::emitNormalizedLoopBounds(RewriterBase &rewriter, Location loc,
+ Value lb, Value ub, Value step) {
+ // For non-index types, generate `arith` instructions
+ // Check if the loop is already known to have a constant zero lower bound or
+ // a constant one step.
+ bool isZeroBased = false;
+ if (auto lbCst = getConstantIntValue(lb))
+ isZeroBased = lbCst.value() == 0;
+
+ bool isStepOne = false;
+ if (auto stepCst = getConstantIntValue(step))
+ isStepOne = stepCst.value() == 1;
+
+ // Compute the number of iterations the loop executes: ceildiv(ub - lb, step)
+ // assuming the step is strictly positive. Update the bounds and the step
+ // of the loop to go from 0 to the number of iterations, if necessary.
+ if (isZeroBased && isStepOne)
+ return {lb, ub, step};
+
+ Value diff =
+ isZeroBased ? ub : rewriter.createOrFold<arith::SubIOp>(loc, ub, lb);
+ Value newUpperBound =
+ isStepOne ? diff
+ : rewriter.createOrFold<arith::CeilDivSIOp>(loc, diff, step);
+
+ Value newLowerBound = isZeroBased
+ ? lb
+ : rewriter.create<arith::ConstantOp>(
+ loc, rewriter.getZeroAttr(lb.getType()));
+ Value newStep = isStepOne
+ ? step
+ : rewriter.create<arith::ConstantOp>(
+ loc, rewriter.getIntegerAttr(step.getType(), 1));
+
+ return {newLowerBound, newUpperBound, newStep};
+}
diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index b51357198b1ca..5454411bc535b 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -245,8 +245,6 @@ void Operation::setOperands(ValueRange operands) {
/// than the range pointed to by 'start'+'length'.
void Operation::setOperands(unsigned start, unsigned length,
ValueRange operands) {
- assert((start + length) <= getNumOperands() &&
- "invalid operand range specified");
if (LLVM_LIKELY(hasOperandStorage))
return getOperandStorage().setOperands(this, start, length, operands);
assert(operands.empty() && "setting operands without an operand storage");
diff --git a/mlir/lib/Interfaces/LoopLikeInterface.cpp b/mlir/lib/Interfaces/LoopLikeInterface.cpp
index 1e0e87b64e811..3f478b9bc0b96 100644
--- a/mlir/lib/Interfaces/LoopLikeInterface.cpp
+++ b/mlir/lib/Interfaces/LoopLikeInterface.cpp
@@ -113,3 +113,28 @@ LogicalResult detail::verifyLoopLikeOpInterface(Operation *op) {
return success();
}
+
+LogicalResult
+detail::verifyLoopLikeWithInductionVarsOpInterface(Operation *op) {
+ auto loopLikeOp = cast<LoopLikeWithInductionVarsOpInterface>(op);
+
+ // Verify number of induction variables, lower bounds, upper bounds and steps.
+ if (loopLikeOp.getInductionVars().size() !=
+ loopLikeOp.getMixedLowerBound().size())
+ return op->emitOpError(
+ "different number of induction variables and lower bounds: ")
+ << loopLikeOp.getInductionVars().size()
+ << " != " << loopLikeOp.getMixedLowerBound().size();
+ if (loopLikeOp.getInductionVars().size() != loopLikeOp.getMixedStep().size())
+ return op->emitOpError(
+ "different number of induction variables and steps: ")
+ << loopLikeOp.getInductionVars().size()
+ << " != " << loopLikeOp.getMixedStep().size();
+ if (loopLikeOp.getInductionVars().size() !=
+ loopLikeOp.getMixedUpperBound().size())
+ return op->emitOpError(
+ "different number of induction variables and upper bounds: ")
+ << loopLikeOp.getInductionVars().size()
+ << " != " << loopLikeOp.getMixedUpperBound().size();
+ return success();
+}
diff --git a/mlir/lib/Transforms/CMakeLists.txt b/mlir/lib/Transforms/CMakeLists.txt
index 90c0298fb5e46..dc24da367ca48 100644
--- a/mlir/lib/Transforms/CMakeLists.txt
+++ b/mlir/lib/Transforms/CMakeLists.txt
@@ -10,6 +10,7 @@ add_mlir_library(MLIRTransforms
LocationSnapshot.cpp
LoopInvariantCodeMotion.cpp
Mem2Reg.cpp
+ NormalizeLoopBounds.cpp
OpStats.cpp
PrintIR.cpp
RemoveDeadValues.cpp
@@ -30,6 +31,7 @@ add_mlir_library(MLIRTransforms
LINK_LIBS PUBLIC
MLIRAnalysis
MLIRCopyOpInterface
+ MLIRDialectUtils
MLIRFunctionInterfaces
MLIRLoopLikeInterface
MLIRMemorySlotInterfaces
diff --git a/mlir/lib/Transforms/NormalizeLoopBounds.cpp b/mlir/lib/Transforms/NormalizeLoopBounds.cpp
new file mode 100644
index 0000000000000..ffb51b05ce1ca
--- /dev/null
+++ b/mlir/lib/Transforms/NormalizeLoopBounds.cpp
@@ -0,0 +1,118 @@
+// Copyright 2024 The IREE Authors
+//
+// Licensed 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/Transforms/Passes.h"
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Utils/LoopUtils.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_NORMALIZELOOPBOUNDS
+#include "mlir/Transforms/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+
+/// Normalize a loop-like operation with induction variables, i.e. calculate
+/// new normalized upper bounds for lower bounds equal to zero and step sizes
+/// equal to one. Then, insert new `affine.apply` operations to calculate the
+/// denormalized index values and update all usage from the original induction
+/// variables to the results of the `affine.apply` operations.
+///
+/// Example:
+/// Transform a `scf.forall` loop with a strictly positive steps
+/// forall (%i, %j) = (%lb0, %lb1) to (%ub0, %ub1) step (%s0, %s1)
+/// into a 0-based loop with step 1
+/// forall (%i, %j) in (ceildiv(%ub0 - %lb0, %s0), ceildiv(%ub1 - %lb1, %s1))
+LogicalResult
+normalizeLoopBounds(RewriterBase &rewriter,
+ LoopLikeWithInductionVarsOpInterface loopLikeOp) {
+ OpBuilder::InsertionGuard g(rewriter);
+ if (loopLikeOp.isNormalized())
+ return success();
+
+ SmallVector<Value> newLbs;
+ SmallVector<Value> newUbs;
+ SmallVector<Value> newSteps;
+ rewriter.setInsertionPoint(loopLikeOp);
+ for (auto &&[iv, lb, ub, step] : llvm::zip(
+ loopLikeOp.getInductionVars(), loopLikeOp.getLowerBound(rewriter),
+ loopLikeOp.getUpperBound(rewriter), loopLikeOp.getStep(rewriter))) {
+ std::optional<int64_t> lbInt = getConstantIntValue(lb);
+ std::optional<int64_t> stepInt = getConstantIntValue(step);
+
+ rewriter.setInsertionPoint(loopLikeOp);
+ auto newLoopParams =
+ emitNormalizedLoopBounds(rewriter, loopLikeOp.getLoc(), lb, ub, step);
+
+ newLbs.push_back(newLoopParams.lowerBound);
+ newUbs.push_back(newLoopParams.upperBound);
+ newSteps.push_back(newLoopParams.step);
+
+ Region ®ion = loopLikeOp.getOperation()->getRegion(0);
+ rewriter.setInsertionPointToStart(®ion.front());
+ SmallVector<Value> operands = {iv};
+ AffineExpr idxExpr, stepExpr, offsetExpr, res;
+ if (!lbInt && !stepInt) {
+ bindDims(loopLikeOp.getContext(), idxExpr, stepExpr, offsetExpr);
+ res = idxExpr * stepExpr + offsetExpr;
+ operands.push_back(step);
+ operands.push_back(lb);
+ } else if (!lbInt) {
+ bindDims(loopLikeOp.getContext(), idxExpr, offsetExpr);
+ res = idxExpr * stepInt.value() + offsetExpr;
+ operands.push_back(lb);
+ } else if (!stepInt) {
+ bindDims(loopLikeOp.getContext(), idxExpr, stepExpr);
+ res = idxExpr * stepExpr + lbInt.value();
+ operands.push_back(step);
+ } else {
+ bindDims(loopLikeOp.getContext(), idxExpr);
+ res = idxExpr * stepInt.value() + lbInt.value();
+ }
+
+ auto affineApply = rewriter.create<affine::AffineApplyOp>(
+ loopLikeOp.getLoc(), res, operands);
+ SmallPtrSet<Operation *, 2> preserve(
+ {iv.getDefiningOp(), affineApply.getOperation()});
+ rewriter.replaceAllUsesExcept(iv, affineApply.getResult(), preserve);
+ }
+
+ rewriter.setInsertionPoint(loopLikeOp);
+ rewriter.modifyOpInPlace(loopLikeOp, [&]() {
+ loopLikeOp.setLowerBounds(newLbs);
+ loopLikeOp.setUpperBounds(newUbs);
+ loopLikeOp.setSteps(newSteps);
+ });
+ return success();
+}
+
+namespace {
+
+/// Pass which normalizes the loop bounds of operations implementing
+/// `LoopLikeWithInductionVarsOpInterface`.
+struct NormalizeLoopBounds
+ : public impl::NormalizeLoopBoundsBase<NormalizeLoopBounds> {
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<affine::AffineDialect>();
+ }
+
+ void runOnOperation() override {
+ Operation *parentOp = getOperation();
+ IRRewriter rewriter(parentOp->getContext());
+
+ parentOp->walk([&](LoopLikeWithInductionVarsOpInterface loopLikeOp) {
+ (void)normalizeLoopBounds(rewriter, loopLikeOp);
+ });
+ }
+};
+
+} // namespace
+
+std::unique_ptr<Pass> mlir::createNormalizeLoopBoundsPass() {
+ return std::make_unique<NormalizeLoopBounds>();
+}
diff --git a/mlir/test/Dialect/Affine/loop-coalescing.mlir b/mlir/test/Dialect/Affine/loop-coalescing.mlir
index ae0adf5a0a02d..0a96e01162d48 100644
--- a/mlir/test/Dialect/Affine/loop-coalescing.mlir
+++ b/mlir/test/Dialect/Affine/loop-coalescing.mlir
@@ -72,19 +72,16 @@ func.func @multi_use() {
return
}
-func.func @unnormalized_loops() {
+// CHECK: %[[orig_ub_i:.*]]: index, %[[orig_ub_j:.*]]: index
+func.func @unnormalized_loops(%ubi: index, %ubj: index) {
// CHECK: %[[orig_step_i:.*]] = arith.constant 2
// CHECK: %[[orig_step_j:.*]] = arith.constant 3
// CHECK: %[[orig_lb_i:.*]] = arith.constant 5
// CHECK: %[[orig_lb_j:.*]] = arith.constant 7
- // CHECK: %[[orig_ub_i:.*]] = arith.constant 10
- // CHECK: %[[orig_ub_j:.*]] = arith.constant 17
%c2 = arith.constant 2 : index
%c3 = arith.constant 3 : index
%c5 = arith.constant 5 : index
%c7 = arith.constant 7 : index
- %c10 = arith.constant 10 : index
- %c17 = arith.constant 17 : index
// Number of iterations in the outer scf.
// CHECK: %[[diff_i:.*]] = arith.subi %[[orig_ub_i]], %[[orig_lb_i]]
@@ -101,10 +98,10 @@ func.func @unnormalized_loops() {
// New bounds of the outer scf.
// CHECK: %[[range:.*]] = arith.muli %[[numiter_i]], %[[numiter_j]]
// CHECK: scf.for %[[i:.*]] = %[[lb_i]] to %[[range]] step %[[step_i]]
- scf.for %i = %c5 to %c10 step %c2 {
+ scf.for %i = %c5 to %ubi step %c2 {
// The inner loop has been removed.
// CHECK-NOT: scf.for
- scf.for %j = %c7 to %c17 step %c3 {
+ scf.for %j = %c7 to %ubj step %c3 {
// The IVs are rewritten.
// CHECK: %[[normalized_j:.*]] = arith.remsi %[[i]], %[[numiter_j]]
// CHECK: %[[normalized_i:.*]] = arith.divsi %[[i]], %[[numiter_j]]
diff --git a/mlir/test/Transforms/normalize-loop-bounds.mlir b/mlir/test/Transforms/normalize-loop-bounds.mlir
new file mode 100644
index 0000000000000..5130f4282b36b
--- /dev/null
+++ b/mlir/test/Transforms/normalize-loop-bounds.mlir
@@ -0,0 +1,266 @@
+// RUN: mlir-opt %s -split-input-file -normalize-loop-bounds -verify-diagnostics | FileCheck %s
+
+// CHECK: #[[$MAP:.+]] = affine_map<(d0) -> (d0 + 2)>
+// CHECK-LABEL: func.func @for_lowerbound_static
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C6:.+]] = arith.constant 6 : index
+// CHECK: scf.for %[[ARG:.+]] = %[[C0]] to %[[C6]] step %[[C1]]
+// CHECK-NEXT: affine.apply #[[$MAP]](%[[ARG]])
+module {
+ func.func @for_lowerbound_static() {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c8 = arith.constant 8 : index
+ scf.for %arg0 = %c2 to %c8 step %c1 {
+ }
+ return
+ }
+}
+
+// -----
+
+// CHECK: #[[$MAP:.+]] = affine_map<(d0, d1) -> (d0 + d1)>
+// CHECK-LABEL: func.func @for_lowerbound_dynamic
+// CHECK-SAME: %[[ARG0:.+]]: index
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index
+// CHECK-DAG: %[[UB:.+]] = arith.subi %[[C8]], %[[ARG0]] : index
+// CHECK: scf.for %[[ARG:.+]] = %[[C0]] to %[[UB]] step %[[C1]]
+// CHECK-NEXT: affine.apply #[[$MAP]](%[[ARG]], %[[ARG0]])
+module {
+ func.func @for_lowerbound_dynamic(%lb: index) {
+ %c1 = arith.constant 1 : index
+ %c8 = arith.constant 8 : index
+ scf.for %arg0 = %lb to %c8 step %c1 {
+ }
+ return
+ }
+}
+
+// -----
+
+// CHECK: #[[$MAP:.+]] = affine_map<(d0) -> (d0 * 2)>
+// CHECK-LABEL: func.func @for_step_static
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C4:.+]] = arith.constant 4 : index
+// CHECK: scf.for %[[ARG:.+]] = %[[C0]] to %[[C4]] step %[[C1]]
+// CHECK-NEXT: affine.apply #[[$MAP]](%[[ARG]])
+module {
+ func.func @for_step_static() {
+ %c0 = arith.constant 0 : index
+ %c2 = arith.constant 2 : index
+ %c8 = arith.constant 8 : index
+ scf.for %arg0 = %c0 to %c8 step %c2 {
+ }
+ return
+ }
+}
+
+// -----
+
+// CHECK: #[[$MAP:.+]] = affine_map<(d0, d1) -> (d0 * d1)>
+// CHECK-LABEL: func.func @for_step_dynamic
+// CHECK-SAME: %[[ARG0:.+]]: index
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index
+// CHECK-DAG: %[[UB:.+]] = arith.ceildivsi %[[C8]], %[[ARG0]] : index
+// CHECK: scf.for %[[ARG:.+]] = %[[C0]] to %[[UB]] step %[[C1]]
+// CHECK-NEXT: affine.apply #[[$MAP]](%[[ARG]], %[[ARG0]])
+module {
+ func.func @for_step_dynamic(%step: index) {
+ %c0 = arith.constant 0 : index
+ %c8 = arith.constant 8 : index
+ scf.for %arg0 = %c0 to %c8 step %step {
+ }
+ return
+ }
+}
+
+// -----
+
+// CHECK: #[[$MAP:.+]] = affine_map<(d0) -> (d0 * 4 + 1)>
+// CHECK-LABEL: func.func @for_lowerbound_and_step_static
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C3:.+]] = arith.constant 3 : index
+// CHECK: scf.for %[[ARG:.+]] = %[[C0]] to %[[C3]] step %[[C1]]
+// CHECK-NEXT: affine.apply #[[$MAP]](%[[ARG]])
+module {
+ func.func @for_lowerbound_and_step_static() {
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c13 = arith.constant 13 : index
+ scf.for %arg0 = %c1 to %c13 step %c4 {
+ }
+ return
+ }
+}
+
+// -----
+
+// CHECK: #[[$MAP:.+]] = affine_map<(d0, d1, d2) -> (d0 * d1 + d2)>
+// CHECK-LABEL: func.func @for_lowerbound_and_step_dynamic
+// CHECK-SAME: %[[LB:.+]]: index, %[[STEP:.+]]: index
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C13:.+]] = arith.constant 13 : index
+// CHECK-DAG: %[[SUB:.+]] = arith.subi %[[C13]], %[[LB]] : index
+// CHECK-DAG: %[[UB:.+]] = arith.ceildivsi %[[SUB]], %[[STEP]] : index
+// CHECK: scf.for %[[ARG:.+]] = %[[C0]] to %[[UB]] step %[[C1]]
+// CHECK-NEXT: affine.apply #[[$MAP]](%[[ARG]], %[[STEP]], %[[LB]])
+module {
+ func.func @for_lowerbound_and_step_dynamic(%lb: index, %step: index) {
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c13 = arith.constant 13 : index
+ scf.for %arg0 = %lb to %c13 step %step {
+ }
+ return
+ }
+}
+
+// -----
+
+// CHECK-DAG: #[[$MAP0:.+]] = affine_map<(d0) -> (d0 + 4)>
+// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0) -> (d0 + 2)>
+// CHECK-LABEL: func.func @forall_lowerbound_static
+// CHECK: scf.forall (%[[ARG0:.+]], %[[ARG1:.+]]) in (6, 12)
+// CHECK-DAG: affine.apply #[[$MAP1]](%[[ARG0]])
+// CHECK-DAG: affine.apply #[[$MAP0]](%[[ARG1]])
+module {
+ func.func @forall_lowerbound_static() {
+ scf.forall (%arg2, %arg3) = (2, 4) to (8, 16) step (1, 1) {
+ }
+ return
+ }
+}
+
+// -----
+
+// CHECK-DAG: #[[$MAP0:.+]] = affine_map<(d0, d1) -> (d0 + d1)>
+// CHECK-LABEL: func.func @forall_lowerbound_dynamic
+// CHECK-SAME: %[[LB0:.+]]: index, %[[LB1:.+]]: index
+// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[UB0:.+]] = arith.subi %[[C8]], %[[LB0]] : index
+// CHECK-DAG: %[[UB1:.+]] = arith.subi %[[C16]], %[[LB1]] : index
+// CHECK: scf.forall (%[[ARG0:.+]], %[[ARG1:.+]]) in (%[[UB0]], %[[UB1]])
+// CHECK-DAG: affine.apply #[[$MAP0]](%[[ARG0]], %[[LB0]])
+// CHECK-DAG: affine.apply #[[$MAP0]](%[[ARG1]], %[[LB1]])
+module {
+ func.func @forall_lowerbound_dynamic(%lb0: index, %lb1: index) {
+ scf.forall (%arg2, %arg3) = (%lb0, %lb1) to (8, 16) step (1, 1) {
+ }
+ return
+ }
+}
+
+// -----
+
+// CHECK: #[[$MAP:.+]] = affine_map<(d0) -> (d0 * 8)>
+// CHECK-LABEL: func.func @forall_step_static
+// CHECK: scf.forall (%[[ARG0:.+]], %[[ARG1:.+]]) in (1, 2)
+// CHECK-DAG: affine.apply #[[$MAP]](%[[ARG0]])
+// CHECK-DAG: affine.apply #[[$MAP]](%[[ARG1]])
+module {
+ func.func @forall_step_static() {
+ scf.forall (%arg2, %arg3) = (0, 0) to (8, 16) step (8, 8) {
+ }
+ return
+ }
+}
+
+// -----
+
+// CHECK: #[[$MAP:.+]] = affine_map<(d0, d1) -> (d0 * d1)>
+// CHECK-LABEL: func.func @forall_step_dynamic
+// CHECK-SAME: %[[STEP0:.+]]: index, %[[STEP1:.+]]: index
+// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[UB0:.+]] = arith.ceildivsi %[[C8]], %[[STEP0]] : index
+// CHECK-DAG: %[[UB1:.+]] = arith.ceildivsi %[[C16]], %[[STEP1]] : index
+// CHECK: scf.forall (%[[ARG0:.+]], %[[ARG1:.+]]) in (%[[UB0]], %[[UB1]])
+// CHECK-DAG: affine.apply #[[$MAP]](%[[ARG0]], %[[STEP0]])
+// CHECK-DAG: affine.apply #[[$MAP]](%[[ARG1]], %[[STEP1]])
+module {
+ func.func @forall_step_dynamic(%step0: index, %step1: index) {
+ scf.forall (%arg2, %arg3) = (0, 0) to (8, 16) step (%step0, %step1) {
+ }
+ return
+ }
+}
+
+// -----
+
+// CHECK-DAG: #[[$MAP0:.+]] = affine_map<(d0) -> (d0 * 4 + 4)>
+// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0) -> (d0 * 2 + 2)>
+// CHECK-LABEL: func.func @forall_lowerbound_and_step_static
+// CHECK: scf.forall (%[[ARG0:.+]], %[[ARG1:.+]]) in (3, 3)
+// CHECK-DAG: affine.apply #[[$MAP1]](%[[ARG0]])
+// CHECK-DAG: affine.apply #[[$MAP0]](%[[ARG1]])
+module {
+ func.func @forall_lowerbound_and_step_static() {
+ scf.forall (%arg2, %arg3) = (2, 4) to (8, 16) step (2, 4) {
+ }
+ return
+ }
+}
+
+// -----
+
+// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2) -> (d0 * d1 + d2)>
+// CHECK-LABEL: func.func @forall_lowerbound_and_step_dynamic
+// CHECK-SAME: %[[LB0:.+]]: index, %[[LB1:.+]]: index, %[[STEP0:.+]]: index, %[[STEP1:.+]]: index
+// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[SUB0:.+]] = arith.subi %[[C8]], %[[LB0]] : index
+// CHECK-DAG: %[[SUB1:.+]] = arith.subi %[[C16]], %[[LB1]] : index
+// CHECK-DAG: %[[UB0:.+]] = arith.ceildivsi %[[SUB0]], %[[STEP0]] : index
+// CHECK-DAG: %[[UB1:.+]] = arith.ceildivsi %[[SUB1]], %[[STEP1]] : index
+// CHECK: scf.forall (%[[ARG0:.+]], %[[ARG1:.+]]) in (%[[UB0]], %[[UB1]])
+// CHECK-DAG: affine.apply #[[$MAP]](%[[ARG0]], %[[STEP0]], %[[LB0]])
+// CHECK-DAG: affine.apply #[[$MAP]](%[[ARG1]], %[[STEP1]], %[[LB1]])
+module {
+ func.func @forall_lowerbound_and_step_dynamic(%lb0: index, %lb1: index, %step0: index, %step1: index) {
+ scf.forall (%arg2, %arg3) = (%lb0, %lb1) to (8, 16) step (%step0, %step1) {
+ }
+ return
+ }
+}
+
+// -----
+
+// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0) -> (d0 * 4 + 2)>
+// CHECK-LABEL: func.func @forall_with_shared_outs_static
+// CHECK-SAME: %[[OUT:.+]]: tensor<200x100xf32>
+// CHECK: scf.forall (%[[ARG0:.+]]) in (2) shared_outs(%{{.+}} = %[[OUT]])
+// CHECK-DAG: affine.apply #[[$MAP]](%[[ARG0]])
+module {
+ func.func @forall_with_shared_outs_static(%out: tensor<200x100xf32>) {
+ scf.forall (%arg0) = (2) to (8) step (4) shared_outs (%o = %out) -> tensor<200x100xf32> {
+ }
+ return
+ }
+}
+
+// -----
+
+// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2) -> (d0 * d1 + d2)>
+// CHECK-LABEL: func.func @forall_with_shared_outs_dynamic
+// CHECK-SAME: %[[LB:.+]]: index, %[[STEP:.+]]: index, %[[OUT:.+]]: tensor<200x100xf32>
+// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index
+// CHECK-DAG: %[[SUB:.+]] = arith.subi %[[C8]], %[[LB]] : index
+// CHECK-DAG: %[[UB:.+]] = arith.ceildivsi %[[SUB]], %[[STEP]] : index
+// CHECK: scf.forall (%[[ARG:.+]]) in (%[[UB]]) shared_outs(%{{.+}} = %[[OUT]])
+// CHECK-DAG: affine.apply #[[$MAP]](%[[ARG]], %[[STEP]], %[[LB]])
+module {
+ func.func @forall_with_shared_outs_dynamic(%lb: index, %step: index, %out: tensor<200x100xf32>) {
+ scf.forall (%arg0) = (%lb) to (8) step (%step) shared_outs (%o = %out) -> tensor<200x100xf32> {
+ }
+ return
+ }
+}
More information about the Mlir-commits
mailing list