Skip to content

Commit

Permalink
Enable specializing both symbol and dim inputs of affine.apply into f…
Browse files Browse the repository at this point in the history 
…or loops
  • Loading branch information
@erwei-xilinx
erwei-xilinx committed Nov 27, 2024
1 parent ab20df5 commit 873ebe8
Show file tree
Hide file tree
Showing 5 changed files with 63 additions and 29 deletions.
7 changes: 3 additions & 4 deletions mlir/include/air/Util/Util.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -225,10 +225,9 @@ std::optional<int> getOffsetDimFromMemrefDim(int dimOnMemref,

// Evaluate the affine expression of affine map on a sparse vector of constant
// ints.
std::optional<int64_t>
evaluateConstantsInMap(AffineMap map,
SmallVector<std::optional<int64_t>> const_inputs,
MLIRContext *ctx);
std::optional<int64_t> evaluateConstantsInMap(
AffineMap map, SmallVector<std::optional<int64_t>> symbolInputs,
SmallVector<std::optional<int64_t>> dimInputs, MLIRContext *ctx);

// Extend the lookupOrDefault method to operate on a vector of values.
Value lookupOrDefaultRange(Value v, IRMapping &remap);
Expand Down
8 changes: 4 additions & 4 deletions mlir/lib/Transform/AIRDependencyScheduleOpt.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1673,12 +1673,12 @@ struct CanonicalizeAffineApplyOnLoopInductionVar
apply.getAffineMap(),
SmallVector<std::optional<int64_t>>{
*mlir::getConstantIntValue(sfo.getUpperBound())},
ctx);
SmallVector<std::optional<int64_t>>{}, ctx);
auto new_lb = air::evaluateConstantsInMap(
apply.getAffineMap(),
SmallVector<std::optional<int64_t>>{
*mlir::getConstantIntValue(sfo.getLowerBound())},
ctx);
SmallVector<std::optional<int64_t>>{}, ctx);
assert(new_ub && new_lb);
int newStepInInt = llvm::divideCeilSigned(*new_ub - *new_lb, tripCount);
IRMapping remap;
Expand Down Expand Up @@ -1706,11 +1706,11 @@ struct CanonicalizeAffineApplyOnLoopInductionVar
auto new_ub = air::evaluateConstantsInMap(
apply.getAffineMap(),
SmallVector<std::optional<int64_t>>{afo.getConstantUpperBound()},
ctx);
SmallVector<std::optional<int64_t>>{}, ctx);
auto new_lb = air::evaluateConstantsInMap(
apply.getAffineMap(),
SmallVector<std::optional<int64_t>>{afo.getConstantLowerBound()},
ctx);
SmallVector<std::optional<int64_t>>{}, ctx);
assert(new_ub && new_lb);
int newStepInInt = llvm::divideCeilSigned(*new_ub - *new_lb, tripCount);
IRMapping remap;
Expand Down
17 changes: 12 additions & 5 deletions mlir/lib/Transform/AIRMiscPasses.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1069,15 +1069,22 @@ void AIRSplitL2MemrefForBufferConstraintPass::partitionMemref(
for (auto user : iv.getUsers())
oneUser = user;
if (auto apply = dyn_cast<affine::AffineApplyOp>(oneUser)) {
SmallVector<std::optional<int64_t>> const_ints;
for (auto oper : apply->getOperands()) {
SmallVector<std::optional<int64_t>> sym_ints;
SmallVector<std::optional<int64_t>> dim_ints;
for (auto oper : apply.getSymbolOperands()) {
if (auto constVal = getConstantIntValue(oper))
const_ints.push_back(constVal);
sym_ints.push_back(constVal);
else
const_ints.push_back(lb);
sym_ints.push_back(lb);
}
for (auto oper : apply.getDimOperands()) {
if (auto constVal = getConstantIntValue(oper))
dim_ints.push_back(constVal);
else
dim_ints.push_back(lb);
}
auto key_opt = air::evaluateConstantsInMap(apply.getAffineMap(),
const_ints, ctx);
sym_ints, dim_ints, ctx);
if (!key_opt)
return;
offset_key = *key_opt;
Expand Down
48 changes: 32 additions & 16 deletions mlir/lib/Util/Util.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1050,9 +1050,11 @@ LogicalResult air::foldForLoopNestAsExtendedSizesAndStrides(

// Evaluate offset from affine map.
auto evalOffsetFromAffineMap = [&](MLIRContext *ctx, AffineMap map) {
SmallVector<std::optional<int64_t>> zeros(map.getNumSymbols(),
std::optional<int64_t>{0});
return air::evaluateConstantsInMap(map, zeros, ctx);
SmallVector<std::optional<int64_t>> zeroSyms(map.getNumSymbols(),
std::optional<int64_t>{0});
SmallVector<std::optional<int64_t>> zeroDims(map.getNumDims(),
std::optional<int64_t>{0});
return air::evaluateConstantsInMap(map, zeroSyms, zeroDims, ctx);
};
for (auto o : for_loops) {
int64_t stepSize = -1;
Expand Down Expand Up @@ -1086,18 +1088,25 @@ LogicalResult air::foldForLoopNestAsExtendedSizesAndStrides(
if (auto exec = dyn_cast<air::ExecuteOp>(iv_consumer))
iv_consumer = &exec.getChildOps().front();
if (auto affop = dyn_cast<affine::AffineApplyOp>(iv_consumer)) {
auto idx = llvm::find_if(affop.getSymbolOperands(),
auto idx = llvm::find_if(affop.getOperands(),
[iv](Value oper) { return oper == iv; });
if (idx != affop.getSymbolOperands().end()) {
if (idx != affop.getOperands().end()) {
auto map = affop.getAffineMap();
int64_t map_offset =
evalOffsetFromAffineMap(for_op->getContext(), map).value();
ind_var_factor = *getConstantIntValue(strides[i]);
SmallVector<std::optional<int64_t>> stepSizeAsVec(
affop.getSymbolOperands().size(), std::optional<int64_t>{0});
stepSizeAsVec[idx - affop.getSymbolOperands().begin()] = stepSize;
SmallVector<std::optional<int64_t>> stepSizeAsSymVec(
affop.getMap().getNumSymbols(), std::optional<int64_t>{0});
SmallVector<std::optional<int64_t>> stepSizeAsDimVec(
affop.getMap().getNumDims(), std::optional<int64_t>{0});
if (idx - affop.getOperands().begin() < affop.getMap().getNumDims())
stepSizeAsDimVec[idx - affop.getOperands().begin()] = stepSize;
else
stepSizeAsSymVec[idx - affop.getOperands().begin() -
affop.getMap().getNumDims()] = stepSize;
int64_t map_gradient = air::evaluateConstantsInMap(
map, stepSizeAsVec, for_op->getContext())
map, stepSizeAsSymVec, stepSizeAsDimVec,
for_op->getContext())
.value() -
map_offset;
ind_var_factor *= map_gradient;
Expand Down Expand Up @@ -1573,21 +1582,28 @@ air::getOffsetDimFromMemrefDim(int dimOnMemref, SmallVector<Value> strides,

// Evaluate the affine expression of affine map on a sparse vector of constant
// ints.
std::optional<int64_t>
air::evaluateConstantsInMap(AffineMap map,
SmallVector<std::optional<int64_t>> const_inputs,
MLIRContext *ctx) {
std::optional<int64_t> air::evaluateConstantsInMap(
AffineMap map, SmallVector<std::optional<int64_t>> symbolInputs,
SmallVector<std::optional<int64_t>> dimInputs, MLIRContext *ctx) {
std::optional<int64_t> output = std::nullopt;
if (map.getNumInputs() != const_inputs.size())
if (map.getNumSymbols() != symbolInputs.size())
return output;
if (map.getNumDims() != dimInputs.size())
return output;
auto newmap = map;
for (unsigned i = 0; i < map.getNumSymbols(); i++) {
if (!const_inputs[i])
if (!symbolInputs[i])
continue;
auto c = getAffineConstantExpr(*const_inputs[i], ctx);
auto c = getAffineConstantExpr(*symbolInputs[i], ctx);
newmap =
newmap.replace(getAffineSymbolExpr(i, ctx), c, 0, map.getNumSymbols());
}
for (unsigned i = 0; i < map.getNumDims(); i++) {
if (!dimInputs[i])
continue;
auto c = getAffineConstantExpr(*dimInputs[i], ctx);
newmap = newmap.replace(getAffineDimExpr(i, ctx), c, map.getNumDims(), 0);
}
output = simplifyAffineMap(newmap).getSingleConstantResult();
return output;
}
Expand Down
12 changes: 12 additions & 0 deletions mlir/test/Transform/AIRDependencyScheduleOpt/specialize-channel-wrap-and-stride.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -386,6 +386,7 @@ module {
// Affine.apply with map joining two for loops in a loop nest.
// CHECK-LABEL: test11

// CHECK: air.channel.put async {{.*}}@channel_26[%c0{{.*}}, %c0{{.*}}] (%{{.*}}[%c0{{.*}}, %c0{{.*}}, %c0{{.*}}] [%c4{{.*}}, %c18{{.*}}, %c4{{.*}}] [%c96{{.*}}, %c16{{.*}}, %c1{{.*}}]) : (memref<1x6x6x16xbf16, 1>)
// CHECK: air.channel.put async {{.*}}@channel_26[%c0{{.*}}, %c0{{.*}}] (%{{.*}}[%c0{{.*}}, %c0{{.*}}, %c0{{.*}}] [%c4{{.*}}, %c18{{.*}}, %c4{{.*}}] [%c96{{.*}}, %c16{{.*}}, %c1{{.*}}]) : (memref<1x6x6x16xbf16, 1>)
// CHECK: air.channel.put async {{.*}}@channel_26[%c0{{.*}}, %c0{{.*}}] (%{{.*}}[%c0{{.*}}, %c0{{.*}}, %c0{{.*}}, %c12{{.*}}] [%c3{{.*}}, %c3{{.*}}, %c4{{.*}}, %c4{{.*}}] [%c96{{.*}}, %c16{{.*}}, %c16{{.*}}, %c1{{.*}}]) : (memref<1x3x6x16xi32, 1>)

Expand Down Expand Up @@ -423,6 +424,17 @@ module {
}
scf.yield %2 : !air.async.token
}
%5 = scf.for %arg9 = %c0 to %c4_1 step %c1 iter_args(%arg13 = %async_token) -> (!air.async.token) {
%2 = scf.for %arg10 = %c0 to %c3_0 step %c1 iter_args(%arg11 = %arg13) -> (!air.async.token) {
%async_token_2, %results_3 = air.execute [%arg11] -> (index) {
%6 = affine.apply #map1()[%arg9, %arg10]
air.execute_terminator %6 : index
}
%3 = air.channel.put async [%async_token_2] @channel_26[%c0, %c0] (%results[%c0, %results_3, %c0, %c0] [%c1, %c1, %c6, %c4_1] [%c576, %c96, %c16, %c1]) : (memref<1x6x6x16xbf16, 1>)
scf.yield %3 : !air.async.token
}
scf.yield %2 : !air.async.token
}
scf.for %arg9 = %c0 to %c3_0 step %c1 {
%60 = scf.for %arg10 = %c0 to %c3_0 step %c1 iter_args(%arg13 = %async_token) -> (!air.async.token) {
%async_token_54, %results_55 = air.execute [%arg13] -> (index) {
Expand Down

0 comments on commit 873ebe8

Please sign in to comment.