diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h index e620067c15be9..0eeebda5f1a14 100644 --- a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h +++ b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h @@ -148,6 +148,22 @@ void denormalizeInductionVariable(RewriterBase &rewriter, Location loc, Value normalizedIv, OpFoldResult origLb, OpFoldResult origStep); +/// For each original loop, the value of the induction variable can be obtained +/// by dividing the induction variable of the linearized loop by the total +/// number of iterations of the loops nested in it modulo the number of +/// iterations in this loop (remove the values related to the outer loops): +/// iv_i = floordiv(iv_linear, product-of-loop-ranges-until-i) mod range_i. +/// Compute these iteratively from the innermost loop by creating a "running +/// quotient" of division by the range. +/// Returns the delinearized induction variables and the preserved users. +std::pair, SmallPtrSet> +delinearizeInductionVariable(RewriterBase &rewriter, Location loc, + Value linearizedIv, ArrayRef ubs); + +/// Helper function to multiply a sequence of values. +Value getProductOfIntsOrIndexes(RewriterBase &rewriter, Location loc, + ArrayRef values); + /// Tile a nest of standard for loops rooted at `rootForOp` by finding such /// parametric tile sizes that the outer loops have a fixed number of iterations /// as defined in `sizes`. diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp index e9471c1dbd0b7..42a27e5ff783d 100644 --- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp +++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp @@ -17,6 +17,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/StaticValueUtils.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/IRMapping.h" #include "mlir/IR/OpDefinition.h" @@ -807,7 +808,7 @@ void mlir::denormalizeInductionVariable(RewriterBase &rewriter, Location loc, static OpFoldResult getProductOfIndexes(RewriterBase &rewriter, Location loc, ArrayRef values) { - assert(!values.empty() && "unexecpted empty array"); + assert(!values.empty() && "unexpected empty array"); AffineExpr s0, s1; bindSymbols(rewriter.getContext(), s0, s1); AffineExpr mul = s0 * s1; @@ -819,9 +820,8 @@ static OpFoldResult getProductOfIndexes(RewriterBase &rewriter, Location loc, return products; } -/// Helper function to multiply a sequence of values. -static Value getProductOfIntsOrIndexes(RewriterBase &rewriter, Location loc, - ArrayRef values) { +Value mlir::getProductOfIntsOrIndexes(RewriterBase &rewriter, Location loc, + ArrayRef values) { assert(!values.empty() && "unexpected empty list"); if (getType(values.front()).isIndex()) { SmallVector ofrs = getAsOpFoldResult(values); @@ -835,7 +835,7 @@ static Value getProductOfIntsOrIndexes(RewriterBase &rewriter, Location loc, continue; if (productOf) productOf = - rewriter.create(loc, productOf.value(), v).getResult(); + rewriter.createOrFold(loc, productOf.value(), v); else productOf = v; } @@ -848,17 +848,9 @@ static Value getProductOfIntsOrIndexes(RewriterBase &rewriter, Location loc, return productOf.value(); } -/// For each original loop, the value of the -/// induction variable can be obtained by dividing the induction variable of -/// the linearized loop by the total number of iterations of the loops nested -/// in it modulo the number of iterations in this loop (remove the values -/// related to the outer loops): -/// iv_i = floordiv(iv_linear, product-of-loop-ranges-until-i) mod range_i. -/// Compute these iteratively from the innermost loop by creating a "running -/// quotient" of division by the range. -static std::pair, SmallPtrSet> -delinearizeInductionVariable(RewriterBase &rewriter, Location loc, - Value linearizedIv, ArrayRef ubs) { +std::pair, SmallPtrSet> +mlir::delinearizeInductionVariable(RewriterBase &rewriter, Location loc, + Value linearizedIv, ArrayRef ubs) { if (linearizedIv.getType().isIndex()) { Operation *delinearizedOp = diff --git a/mlir/test/Dialect/SCF/transform-ops.mlir b/mlir/test/Dialect/SCF/transform-ops.mlir index d9445182769e7..81462885dba4c 100644 --- a/mlir/test/Dialect/SCF/transform-ops.mlir +++ b/mlir/test/Dialect/SCF/transform-ops.mlir @@ -500,13 +500,11 @@ func.func @coalesce_i32_loops() { %1 = arith.constant 128 : i32 %2 = arith.constant 2 : i32 %3 = arith.constant 64 : i32 - // CHECK: %[[VAL_4:.*]] = arith.constant 64 : i32 // CHECK: %[[ZERO:.*]] = arith.constant 0 : i32 // CHECK: %[[ONE:.*]] = arith.constant 1 : i32 - // CHECK: %[[VAL_7:.*]] = arith.constant 32 : i32 // CHECK: %[[VAL_8:.*]] = arith.constant 0 : i32 // CHECK: %[[VAL_9:.*]] = arith.constant 1 : i32 - // CHECK: %[[UB:.*]] = arith.muli %[[VAL_4]], %[[VAL_7]] : i32 + // CHECK: %[[UB:.*]] = arith.constant 2048 : i32 // CHECK: scf.for %[[VAL_11:.*]] = %[[ZERO]] to %[[UB]] step %[[ONE]] : i32 { scf.for %i = %0 to %1 step %2 : i32 { scf.for %j = %0 to %3 step %2 : i32 {