diff --git a/compiler/src/iree/compiler/Codegen/Common/DecomposePackUnPackOps.cpp b/compiler/src/iree/compiler/Codegen/Common/DecomposePackUnPackOps.cpp index 771b73bf9425b..898791cf14bb4 100644 --- a/compiler/src/iree/compiler/Codegen/Common/DecomposePackUnPackOps.cpp +++ b/compiler/src/iree/compiler/Codegen/Common/DecomposePackUnPackOps.cpp @@ -176,7 +176,7 @@ void DecomposePackUnPackOpsPass::runOnOperation() { IRRewriter rewriter(ctx); auto packOptions = scf::SCFTileAndFuseOptions().setTilingOptions( scf::SCFTilingOptions().setTileSizeComputationFunction( - [](OpBuilder &builder, Operation *op) -> SmallVector { + [](OpBuilder &builder, Operation *op) -> SmallVector { auto packOp = cast(op); // Do nothing if any of inner tile sizes is dynamic. @@ -187,9 +187,8 @@ void DecomposePackUnPackOpsPass::runOnOperation() { } int inputRank = packOp.getSourceRank(); - SmallVector tileSizes( - inputRank, - builder.create(packOp.getLoc(), 1)); + SmallVector tileSizes(inputRank, + builder.getIndexAttr(1)); return tileSizes; })); funcOp->walk([&](tensor::PackOp op) { @@ -204,18 +203,15 @@ void DecomposePackUnPackOpsPass::runOnOperation() { auto unpackTilingOptions = scf::SCFTilingOptions().setTileSizeComputationFunction( [](OpBuilder &builder, Operation *op) { - Location loc = op->getLoc(); auto unpackOp = cast(op); int numLoops = unpackOp.getDestRank(); auto dimAndTileMapping = unpackOp.getDimAndTileMapping(); - SmallVector tileSizes; + SmallVector tileSizes; for (int i = 0; i < numLoops; ++i) { if (dimAndTileMapping.count(i)) { - tileSizes.push_back(getValueOrCreateConstantIndexOp( - builder, loc, dimAndTileMapping[i])); + tileSizes.push_back(dimAndTileMapping[i]); } else { - tileSizes.push_back( - builder.create(loc, 1)); + tileSizes.push_back(builder.getIndexAttr(1)); } } return tileSizes; diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUTensorTile.cpp b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUTensorTile.cpp index 5e5cc376f4b15..c2664ff065441 100644 --- a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUTensorTile.cpp +++ b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUTensorTile.cpp @@ -77,8 +77,10 @@ class TileConsumerAndFuseInputProducer final } // Tile the current op and fuse its immediate input operands. + SmallVector tileSizesOfr = + getAsIndexOpFoldResult(rewriter.getContext(), tileSizes); FailureOr tilingResult = - tileConsumerAndFuseInputProducer(rewriter, op, tileSizes); + tileConsumerAndFuseInputProducer(rewriter, op, tileSizesOfr); if (failed(tilingResult)) { return rewriter.notifyMatchFailure(op, "failed to tile consumer"); } @@ -94,7 +96,7 @@ class TileConsumerAndFuseInputProducer final FailureOr tileConsumerAndFuseInputProducer(PatternRewriter &rewriter, TilingInterface consumer, - ArrayRef tileSizes) const { + ArrayRef tileSizes) const { // First tile the current op as the consumer op. auto tilingOptions = scf::SCFTilingOptions().setTileSizes(tileSizes); FailureOr tilingResult = @@ -275,7 +277,8 @@ static LogicalResult tileAndUnrollConv(func::FuncOp funcOp) { for (linalg::ConvolutionOpInterface convOp : convOps) { auto consumerOp = cast(*convOp); IRRewriter rewriter(funcOp.getContext()); - SmallVector tileSizes = getTileSizes(consumerOp, 1); + SmallVector tileSizes = getAsIndexOpFoldResult( + funcOp.getContext(), getTileSizes(consumerOp, 1)); if (tileSizes.empty()) return success(); diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUSplitReduction.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUSplitReduction.cpp index 7a347875de5e2..23fffd2627d87 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUSplitReduction.cpp +++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUSplitReduction.cpp @@ -119,8 +119,9 @@ LogicalResult splitReductionImpl(Operation *op, int64_t size, auto numLoops = linalgOp.getNumLoops(); // 1) Tile to extract a single vector-length array. - SmallVector tileSizesSVFirst(numLoops, 1); - tileSizesSVFirst[numLoops - 1] = 0; + SmallVector tileSizesSVFirst(numLoops, + rewriter.getIndexAttr(1)); + tileSizesSVFirst[numLoops - 1] = rewriter.getIndexAttr(0); auto options = scf::SCFTilingOptions().setTileSizes(tileSizesSVFirst); FailureOr tileResFirst = scf::tileUsingSCFForOp( rewriter, cast(linalgOp.getOperation()), options); @@ -142,10 +143,11 @@ LogicalResult splitReductionImpl(Operation *op, int64_t size, // 3) Tile the first op generated by splitReduction with tile size of 1, // to essentially create a reduction loop. Note that // splitRes->splitLinalgOp.getNumLoops() = numLoops + 1. - SmallVector tileSizesSV(splitRes->splitLinalgOp.getNumLoops(), 0); + SmallVector tileSizesSV(splitRes->splitLinalgOp.getNumLoops(), + rewriter.getIndexAttr(0)); // The reduction happens only in the penultimate dimension, which we now // tile. - tileSizesSV[numLoops - 1] = 1; + tileSizesSV[numLoops - 1] = rewriter.getIndexAttr(1); options = scf::SCFTilingOptions().setTileSizes(tileSizesSV); FailureOr tileRes = scf::tileUsingSCFForOp( rewriter, cast(splitRes->splitLinalgOp.getOperation()), diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTile.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTile.cpp index e8a0cccc8347e..e9671f10cc12d 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTile.cpp +++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTile.cpp @@ -29,17 +29,15 @@ namespace { /// Builds a proper tile sizes vector for the op. /// scf::tileUsingSCFForOp expects the num of tile sizes = num of loops. This /// method returns a proper tile sizes vector for each op during tiling. -static SmallVector buildTileSizesForOp(OpBuilder &b, Operation *op, - ArrayRef tileSizes) { +static SmallVector +buildTileSizesForOp(OpBuilder &b, Operation *op, ArrayRef tileSizes) { auto tilingOp = cast(op); SmallVector newTileSizes(tileSizes); newTileSizes.resize(tilingOp.getLoopIteratorTypes().size(), /*default=*/0); OpBuilder::InsertionGuard guard(b); - return llvm::map_to_vector(newTileSizes, [&](int64_t size) -> Value { - return b.create(tilingOp->getLoc(), size); - }); + return getAsIndexOpFoldResult(b.getContext(), newTileSizes); } /// This pass tiles all the TilingInterface operations. The `tilingLevel` must diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTileAndFuse.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTileAndFuse.cpp index ce8b3718de232..31f3faf45b196 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTileAndFuse.cpp +++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTileAndFuse.cpp @@ -245,6 +245,7 @@ void LLVMCPUTileAndFusePass::runOnOperation() { tileSizes = maybeLoweringConfig.value().getTileSizeVals(tilingLevel); } + IRRewriter rewriter(context); int numLoops = consumerOp.getLoopIteratorTypes().size(); if (numLoops > tileSizes.size()) { tileSizes.append(numLoops - tileSizes.size(), 0); @@ -256,8 +257,9 @@ void LLVMCPUTileAndFusePass::runOnOperation() { return; } - auto options = scf::SCFTilingOptions().setTileSizes(tileSizes); - IRRewriter rewriter(context); + SmallVector tileSizesOfr = + getAsIndexOpFoldResult(rewriter.getContext(), tileSizes); + auto options = scf::SCFTilingOptions().setTileSizes(tileSizesOfr); if (failed(applyTileAndFuse(rewriter, consumerOp, options))) { LLVM_DEBUG(llvm::dbgs() << "----- tile and fuse failed -----\n"); return signalPassFailure(); diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/SPIRVTile.cpp b/compiler/src/iree/compiler/Codegen/SPIRV/SPIRVTile.cpp index 97eac04df9dbd..13a7290f7b028 100644 --- a/compiler/src/iree/compiler/Codegen/SPIRV/SPIRVTile.cpp +++ b/compiler/src/iree/compiler/Codegen/SPIRV/SPIRVTile.cpp @@ -115,11 +115,13 @@ static LogicalResult tileAndDistributeToThreads(linalg::LinalgOp consumerOp, ArrayRef tileSizes) { MLIRContext *context = consumerOp.getContext(); IRRewriter rewriter(context); + SmallVector tileSizesOfr = + getAsIndexOpFoldResult(context, tileSizes); FailureOr tileAndFuseResult = scf::tileConsumerAndFuseProducerGreedilyUsingSCFForOp( rewriter, cast(consumerOp.getOperation()), scf::SCFTileAndFuseOptions().setTilingOptions( - scf::SCFTilingOptions().setTileSizes(tileSizes))); + scf::SCFTilingOptions().setTileSizes(tileSizesOfr))); if (failed(tileAndFuseResult)) { return consumerOp.emitOpError("failed tiling and fusing producers"); @@ -240,7 +242,7 @@ static void concretizePadShape(func::FuncOp funcOp) { /// Tiles one of the convolution output window dimensions with size 1 to prepare /// for downsizing 2-D convolution ops into 1-D ones. static LogicalResult tileAndUnrollConvWindow(func::FuncOp funcOp, - ArrayRef tileSizes) { + ArrayRef tileSizes) { SmallVector convOps; funcOp.walk([&convOps](linalg::ConvolutionOpInterface convOp) { convOps.push_back(convOp); @@ -344,7 +346,8 @@ class SPIRVTilePass final : public SPIRVTileBase { fusePadIntoConsumer(funcOp); - SmallVector windowTileSizes = loweringConfig->getTileSizeVals(3); + SmallVector windowTileSizes = + getAsIndexOpFoldResult(context, loweringConfig->getTileSizeVals(3)); if (failed(tileAndUnrollConvWindow(funcOp, windowTileSizes))) { return signalPassFailure(); }