Skip to content

Commit

Permalink
Fixes for llvm/llvm-project#66566
Browse files Browse the repository at this point in the history
  • Loading branch information
@MaheshRavishankar
MaheshRavishankar committed Sep 18, 2023
1 parent bb813ba commit f812714
Show file tree
Hide file tree
Showing 6 changed files with 31 additions and 27 deletions.
16 changes: 6 additions & 10 deletions compiler/src/iree/compiler/Codegen/Common/DecomposePackUnPackOps.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,7 @@ void DecomposePackUnPackOpsPass::runOnOperation() {
IRRewriter rewriter(ctx);
auto packOptions = scf::SCFTileAndFuseOptions().setTilingOptions(
scf::SCFTilingOptions().setTileSizeComputationFunction(
[](OpBuilder &builder, Operation *op) -> SmallVector<Value> {
[](OpBuilder &builder, Operation *op) -> SmallVector<OpFoldResult> {
auto packOp = cast<tensor::PackOp>(op);

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

int inputRank = packOp.getSourceRank();
SmallVector<Value> tileSizes(
inputRank,
builder.create<arith::ConstantIndexOp>(packOp.getLoc(), 1));
SmallVector<OpFoldResult> tileSizes(inputRank,
builder.getIndexAttr(1));
return tileSizes;
}));
funcOp->walk([&](tensor::PackOp op) {
Expand All @@ -204,18 +203,15 @@ void DecomposePackUnPackOpsPass::runOnOperation() {
auto unpackTilingOptions =
scf::SCFTilingOptions().setTileSizeComputationFunction(
[](OpBuilder &builder, Operation *op) {
Location loc = op->getLoc();
auto unpackOp = cast<tensor::UnPackOp>(op);
int numLoops = unpackOp.getDestRank();
auto dimAndTileMapping = unpackOp.getDimAndTileMapping();
SmallVector<Value> tileSizes;
SmallVector<OpFoldResult> 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<arith::ConstantIndexOp>(loc, 1));
tileSizes.push_back(builder.getIndexAttr(1));
}
}
return tileSizes;
Expand Down
9 changes: 6 additions & 3 deletions compiler/src/iree/compiler/Codegen/Common/GPU/GPUTensorTile.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -77,8 +77,10 @@ class TileConsumerAndFuseInputProducer final
}

// Tile the current op and fuse its immediate input operands.
SmallVector<OpFoldResult> tileSizesOfr =
getAsIndexOpFoldResult(rewriter.getContext(), tileSizes);
FailureOr<scf::SCFTilingResult> tilingResult =
tileConsumerAndFuseInputProducer(rewriter, op, tileSizes);
tileConsumerAndFuseInputProducer(rewriter, op, tileSizesOfr);
if (failed(tilingResult)) {
return rewriter.notifyMatchFailure(op, "failed to tile consumer");
}
Expand All @@ -94,7 +96,7 @@ class TileConsumerAndFuseInputProducer final
FailureOr<scf::SCFTilingResult>
tileConsumerAndFuseInputProducer(PatternRewriter &rewriter,
TilingInterface consumer,
ArrayRef<int64_t> tileSizes) const {
ArrayRef<OpFoldResult> tileSizes) const {
// First tile the current op as the consumer op.
auto tilingOptions = scf::SCFTilingOptions().setTileSizes(tileSizes);
FailureOr<scf::SCFTilingResult> tilingResult =
Expand Down Expand Up @@ -275,7 +277,8 @@ static LogicalResult tileAndUnrollConv(func::FuncOp funcOp) {
for (linalg::ConvolutionOpInterface convOp : convOps) {
auto consumerOp = cast<linalg::LinalgOp>(*convOp);
IRRewriter rewriter(funcOp.getContext());
SmallVector<int64_t> tileSizes = getTileSizes(consumerOp, 1);
SmallVector<OpFoldResult> tileSizes = getAsIndexOpFoldResult(
funcOp.getContext(), getTileSizes(consumerOp, 1));
if (tileSizes.empty())
return success();

Expand Down
10 changes: 6 additions & 4 deletions compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUSplitReduction.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<int64_t> tileSizesSVFirst(numLoops, 1);
tileSizesSVFirst[numLoops - 1] = 0;
SmallVector<OpFoldResult> tileSizesSVFirst(numLoops,
rewriter.getIndexAttr(1));
tileSizesSVFirst[numLoops - 1] = rewriter.getIndexAttr(0);
auto options = scf::SCFTilingOptions().setTileSizes(tileSizesSVFirst);
FailureOr<scf::SCFTilingResult> tileResFirst = scf::tileUsingSCFForOp(
rewriter, cast<TilingInterface>(linalgOp.getOperation()), options);
Expand All @@ -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<int64_t> tileSizesSV(splitRes->splitLinalgOp.getNumLoops(), 0);
SmallVector<OpFoldResult> 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<scf::SCFTilingResult> tileRes = scf::tileUsingSCFForOp(
rewriter, cast<TilingInterface>(splitRes->splitLinalgOp.getOperation()),
Expand Down
8 changes: 3 additions & 5 deletions compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTile.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<Value> buildTileSizesForOp(OpBuilder &b, Operation *op,
ArrayRef<int64_t> tileSizes) {
static SmallVector<OpFoldResult>
buildTileSizesForOp(OpBuilder &b, Operation *op, ArrayRef<int64_t> tileSizes) {
auto tilingOp = cast<TilingInterface>(op);

SmallVector<int64_t> 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<arith::ConstantIndexOp>(tilingOp->getLoc(), size);
});
return getAsIndexOpFoldResult(b.getContext(), newTileSizes);
}

/// This pass tiles all the TilingInterface operations. The `tilingLevel` must
Expand Down
6 changes: 4 additions & 2 deletions compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTileAndFuse.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -256,8 +257,9 @@ void LLVMCPUTileAndFusePass::runOnOperation() {
return;
}

auto options = scf::SCFTilingOptions().setTileSizes(tileSizes);
IRRewriter rewriter(context);
SmallVector<OpFoldResult> 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();
Expand Down
9 changes: 6 additions & 3 deletions compiler/src/iree/compiler/Codegen/SPIRV/SPIRVTile.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -115,11 +115,13 @@ static LogicalResult tileAndDistributeToThreads(linalg::LinalgOp consumerOp,
ArrayRef<int64_t> tileSizes) {
MLIRContext *context = consumerOp.getContext();
IRRewriter rewriter(context);
SmallVector<OpFoldResult> tileSizesOfr =
getAsIndexOpFoldResult(context, tileSizes);
FailureOr<scf::SCFTileAndFuseResult> tileAndFuseResult =
scf::tileConsumerAndFuseProducerGreedilyUsingSCFForOp(
rewriter, cast<TilingInterface>(consumerOp.getOperation()),
scf::SCFTileAndFuseOptions().setTilingOptions(
scf::SCFTilingOptions().setTileSizes(tileSizes)));
scf::SCFTilingOptions().setTileSizes(tileSizesOfr)));

if (failed(tileAndFuseResult)) {
return consumerOp.emitOpError("failed tiling and fusing producers");
Expand Down Expand Up @@ -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<int64_t> tileSizes) {
ArrayRef<OpFoldResult> tileSizes) {
SmallVector<linalg::ConvolutionOpInterface, 1> convOps;
funcOp.walk([&convOps](linalg::ConvolutionOpInterface convOp) {
convOps.push_back(convOp);
Expand Down Expand Up @@ -344,7 +346,8 @@ class SPIRVTilePass final : public SPIRVTileBase<SPIRVTilePass> {

fusePadIntoConsumer(funcOp);

SmallVector<int64_t> windowTileSizes = loweringConfig->getTileSizeVals(3);
SmallVector<OpFoldResult> windowTileSizes =
getAsIndexOpFoldResult(context, loweringConfig->getTileSizeVals(3));
if (failed(tileAndUnrollConvWindow(funcOp, windowTileSizes))) {
return signalPassFailure();
}
Expand Down

0 comments on commit f812714

Please sign in to comment.