Skip to content

Commit e654261

Browse files
MaheshRavishankarMaheshRavishankar
committed
1 parent 4aa8f01 commit e654261

6 files changed

Lines changed: 31 additions & 27 deletions

File tree

‎compiler/src/iree/compiler/Codegen/Common/DecomposePackUnPackOps.cpp‎

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -176,7 +176,7 @@ void DecomposePackUnPackOpsPass::runOnOperation() {
176176
IRRewriter rewriter(ctx);
177177
auto packOptions = scf::SCFTileAndFuseOptions().setTilingOptions(
178178
scf::SCFTilingOptions().setTileSizeComputationFunction(
179-
[](OpBuilder &builder, Operation *op) -> SmallVector<Value> {
179+
[](OpBuilder &builder, Operation *op) -> SmallVector<OpFoldResult> {
180180
auto packOp = cast<tensor::PackOp>(op);
181181

182182
// Do nothing if any of inner tile sizes is dynamic.
@@ -187,9 +187,8 @@ void DecomposePackUnPackOpsPass::runOnOperation() {
187187
}
188188

189189
int inputRank = packOp.getSourceRank();
190-
SmallVector<Value> tileSizes(
191-
inputRank,
192-
builder.create<arith::ConstantIndexOp>(packOp.getLoc(), 1));
190+
SmallVector<OpFoldResult> tileSizes(inputRank,
191+
builder.getIndexAttr(1));
193192
return tileSizes;
194193
}));
195194
funcOp->walk([&](tensor::PackOp op) {
@@ -204,18 +203,15 @@ void DecomposePackUnPackOpsPass::runOnOperation() {
204203
auto unpackTilingOptions =
205204
scf::SCFTilingOptions().setTileSizeComputationFunction(
206205
[](OpBuilder &builder, Operation *op) {
207-
Location loc = op->getLoc();
208206
auto unpackOp = cast<tensor::UnPackOp>(op);
209207
int numLoops = unpackOp.getDestRank();
210208
auto dimAndTileMapping = unpackOp.getDimAndTileMapping();
211-
SmallVector<Value> tileSizes;
209+
SmallVector<OpFoldResult> tileSizes;
212210
for (int i = 0; i < numLoops; ++i) {
213211
if (dimAndTileMapping.count(i)) {
214-
tileSizes.push_back(getValueOrCreateConstantIndexOp(
215-
builder, loc, dimAndTileMapping[i]));
212+
tileSizes.push_back(dimAndTileMapping[i]);
216213
} else {
217-
tileSizes.push_back(
218-
builder.create<arith::ConstantIndexOp>(loc, 1));
214+
tileSizes.push_back(builder.getIndexAttr(1));
219215
}
220216
}
221217
return tileSizes;

‎compiler/src/iree/compiler/Codegen/Common/GPU/GPUTensorTile.cpp‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -77,8 +77,10 @@ class TileConsumerAndFuseInputProducer final
7777
}
7878

7979
// Tile the current op and fuse its immediate input operands.
80+
SmallVector<OpFoldResult> tileSizesOfr =
81+
getAsIndexOpFoldResult(rewriter.getContext(), tileSizes);
8082
FailureOr<scf::SCFTilingResult> tilingResult =
81-
tileConsumerAndFuseInputProducer(rewriter, op, tileSizes);
83+
tileConsumerAndFuseInputProducer(rewriter, op, tileSizesOfr);
8284
if (failed(tilingResult)) {
8385
return rewriter.notifyMatchFailure(op, "failed to tile consumer");
8486
}
@@ -94,7 +96,7 @@ class TileConsumerAndFuseInputProducer final
9496
FailureOr<scf::SCFTilingResult>
9597
tileConsumerAndFuseInputProducer(PatternRewriter &rewriter,
9698
TilingInterface consumer,
97-
ArrayRef<int64_t> tileSizes) const {
99+
ArrayRef<OpFoldResult> tileSizes) const {
98100
// First tile the current op as the consumer op.
99101
auto tilingOptions = scf::SCFTilingOptions().setTileSizes(tileSizes);
100102
FailureOr<scf::SCFTilingResult> tilingResult =
@@ -275,7 +277,8 @@ static LogicalResult tileAndUnrollConv(func::FuncOp funcOp) {
275277
for (linalg::ConvolutionOpInterface convOp : convOps) {
276278
auto consumerOp = cast<linalg::LinalgOp>(*convOp);
277279
IRRewriter rewriter(funcOp.getContext());
278-
SmallVector<int64_t> tileSizes = getTileSizes(consumerOp, 1);
280+
SmallVector<OpFoldResult> tileSizes = getAsIndexOpFoldResult(
281+
funcOp.getContext(), getTileSizes(consumerOp, 1));
279282
if (tileSizes.empty())
280283
return success();
281284

‎compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUSplitReduction.cpp‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -112,8 +112,9 @@ LogicalResult splitReductionImpl(Operation *op, int64_t size,
112112
auto numLoops = linalgOp.getNumLoops();
113113

114114
// 1) Tile to extract a single vector-length array.
115-
SmallVector<int64_t> tileSizesSVFirst(numLoops, 1);
116-
tileSizesSVFirst[numLoops - 1] = 0;
115+
SmallVector<OpFoldResult> tileSizesSVFirst(numLoops,
116+
rewriter.getIndexAttr(1));
117+
tileSizesSVFirst[numLoops - 1] = rewriter.getIndexAttr(0);
117118
auto options = scf::SCFTilingOptions().setTileSizes(tileSizesSVFirst);
118119
FailureOr<scf::SCFTilingResult> tileResFirst = scf::tileUsingSCFForOp(
119120
rewriter, cast<TilingInterface>(linalgOp.getOperation()), options);
@@ -135,10 +136,11 @@ LogicalResult splitReductionImpl(Operation *op, int64_t size,
135136
// 3) Tile the first op generated by splitReduction with tile size of 1,
136137
// to essentially create a reduction loop. Note that
137138
// splitRes->splitLinalgOp.getNumLoops() = numLoops + 1.
138-
SmallVector<int64_t> tileSizesSV(splitRes->splitLinalgOp.getNumLoops(), 0);
139+
SmallVector<OpFoldResult> tileSizesSV(splitRes->splitLinalgOp.getNumLoops(),
140+
rewriter.getIndexAttr(0));
139141
// The reduction happens only in the penultimate dimension, which we now
140142
// tile.
141-
tileSizesSV[numLoops - 1] = 1;
143+
tileSizesSV[numLoops - 1] = rewriter.getIndexAttr(1);
142144
options = scf::SCFTilingOptions().setTileSizes(tileSizesSV);
143145
FailureOr<scf::SCFTilingResult> tileRes = scf::tileUsingSCFForOp(
144146
rewriter, cast<TilingInterface>(splitRes->splitLinalgOp.getOperation()),

‎compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTile.cpp‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -29,17 +29,15 @@ namespace {
2929
/// Builds a proper tile sizes vector for the op.
3030
/// scf::tileUsingSCFForOp expects the num of tile sizes = num of loops. This
3131
/// method returns a proper tile sizes vector for each op during tiling.
32-
static SmallVector<Value> buildTileSizesForOp(OpBuilder &b, Operation *op,
33-
ArrayRef<int64_t> tileSizes) {
32+
static SmallVector<OpFoldResult>
33+
buildTileSizesForOp(OpBuilder &b, Operation *op, ArrayRef<int64_t> tileSizes) {
3434
auto tilingOp = cast<TilingInterface>(op);
3535

3636
SmallVector<int64_t> newTileSizes(tileSizes);
3737
newTileSizes.resize(tilingOp.getLoopIteratorTypes().size(), /*default=*/0);
3838

3939
OpBuilder::InsertionGuard guard(b);
40-
return llvm::map_to_vector(newTileSizes, [&](int64_t size) -> Value {
41-
return b.create<arith::ConstantIndexOp>(tilingOp->getLoc(), size);
42-
});
40+
return getAsIndexOpFoldResult(b.getContext(), newTileSizes);
4341
}
4442

4543
/// This pass tiles all the TilingInterface operations. The `tilingLevel` must

‎compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTileAndFuse.cpp‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -245,6 +245,7 @@ void LLVMCPUTileAndFusePass::runOnOperation() {
245245
tileSizes = maybeLoweringConfig.value().getTileSizeVals(tilingLevel);
246246
}
247247

248+
IRRewriter rewriter(context);
248249
int numLoops = consumerOp.getLoopIteratorTypes().size();
249250
if (numLoops > tileSizes.size()) {
250251
tileSizes.append(numLoops - tileSizes.size(), 0);
@@ -256,8 +257,9 @@ void LLVMCPUTileAndFusePass::runOnOperation() {
256257
return;
257258
}
258259

259-
auto options = scf::SCFTilingOptions().setTileSizes(tileSizes);
260-
IRRewriter rewriter(context);
260+
SmallVector<OpFoldResult> tileSizesOfr =
261+
getAsIndexOpFoldResult(rewriter.getContext(), tileSizes);
262+
auto options = scf::SCFTilingOptions().setTileSizes(tileSizesOfr);
261263
if (failed(applyTileAndFuse(rewriter, consumerOp, options))) {
262264
LLVM_DEBUG(llvm::dbgs() << "----- tile and fuse failed -----\n");
263265
return signalPassFailure();

‎compiler/src/iree/compiler/Codegen/SPIRV/SPIRVTile.cpp‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -115,11 +115,13 @@ static LogicalResult tileAndDistributeToThreads(linalg::LinalgOp consumerOp,
115115
ArrayRef<int64_t> tileSizes) {
116116
MLIRContext *context = consumerOp.getContext();
117117
IRRewriter rewriter(context);
118+
SmallVector<OpFoldResult> tileSizesOfr =
119+
getAsIndexOpFoldResult(context, tileSizes);
118120
FailureOr<scf::SCFTileAndFuseResult> tileAndFuseResult =
119121
scf::tileConsumerAndFuseProducerGreedilyUsingSCFForOp(
120122
rewriter, cast<TilingInterface>(consumerOp.getOperation()),
121123
scf::SCFTileAndFuseOptions().setTilingOptions(
122-
scf::SCFTilingOptions().setTileSizes(tileSizes)));
124+
scf::SCFTilingOptions().setTileSizes(tileSizesOfr)));
123125

124126
if (failed(tileAndFuseResult)) {
125127
return consumerOp.emitOpError("failed tiling and fusing producers");
@@ -240,7 +242,7 @@ static void concretizePadShape(func::FuncOp funcOp) {
240242
/// Tiles one of the convolution output window dimensions with size 1 to prepare
241243
/// for downsizing 2-D convolution ops into 1-D ones.
242244
static LogicalResult tileAndUnrollConvWindow(func::FuncOp funcOp,
243-
ArrayRef<int64_t> tileSizes) {
245+
ArrayRef<OpFoldResult> tileSizes) {
244246
SmallVector<linalg::ConvolutionOpInterface, 1> convOps;
245247
funcOp.walk([&convOps](linalg::ConvolutionOpInterface convOp) {
246248
convOps.push_back(convOp);
@@ -344,7 +346,8 @@ class SPIRVTilePass final : public SPIRVTileBase<SPIRVTilePass> {
344346

345347
fusePadIntoConsumer(funcOp);
346348

347-
SmallVector<int64_t> windowTileSizes = loweringConfig->getTileSizeVals(3);
349+
SmallVector<OpFoldResult> windowTileSizes =
350+
getAsIndexOpFoldResult(context, loweringConfig->getTileSizeVals(3));
348351
if (failed(tileAndUnrollConvWindow(funcOp, windowTileSizes))) {
349352
return signalPassFailure();
350353
}

0 commit comments

Comments
 (0)