Skip to content

[MLIR][SCF] Expose scf util functions in header file #138784

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

[MLIR][SCF] Expose scf util functions in header file #138784

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions mlir/include/mlir/Dialect/SCF/Utils/Utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<SmallVector<Value>, SmallPtrSet<Operation *, 2>>
delinearizeInductionVariable(RewriterBase &rewriter, Location loc,
Value linearizedIv, ArrayRef<Value> ubs);

/// Helper function to multiply a sequence of values.
Value getProductOfIntsOrIndexes(RewriterBase &rewriter, Location loc,
ArrayRef<Value> 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`.
Expand Down
24 changes: 8 additions & 16 deletions mlir/lib/Dialect/SCF/Utils/Utils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -807,7 +808,7 @@ void mlir::denormalizeInductionVariable(RewriterBase &rewriter, Location loc,

static OpFoldResult getProductOfIndexes(RewriterBase &rewriter, Location loc,
ArrayRef<OpFoldResult> 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;
Expand All @@ -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<Value> values) {
Value mlir::getProductOfIntsOrIndexes(RewriterBase &rewriter, Location loc,
ArrayRef<Value> values) {
assert(!values.empty() && "unexpected empty list");
if (getType(values.front()).isIndex()) {
SmallVector<OpFoldResult> ofrs = getAsOpFoldResult(values);
Expand All @@ -835,7 +835,7 @@ static Value getProductOfIntsOrIndexes(RewriterBase &rewriter, Location loc,
continue;
if (productOf)
productOf =
rewriter.create<arith::MulIOp>(loc, productOf.value(), v).getResult();
rewriter.createOrFold<arith::MulIOp>(loc, productOf.value(), v);
else
productOf = v;
}
Expand All @@ -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<SmallVector<Value>, SmallPtrSet<Operation *, 2>>
delinearizeInductionVariable(RewriterBase &rewriter, Location loc,
Value linearizedIv, ArrayRef<Value> ubs) {
std::pair<SmallVector<Value>, SmallPtrSet<Operation *, 2>>
mlir::delinearizeInductionVariable(RewriterBase &rewriter, Location loc,
Value linearizedIv, ArrayRef<Value> ubs) {

if (linearizedIv.getType().isIndex()) {
Operation *delinearizedOp =
Expand Down
4 changes: 1 addition & 3 deletions mlir/test/Dialect/SCF/transform-ops.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down