diff --git a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h --- a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h +++ b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h @@ -127,6 +127,11 @@ ValueRange ivs, ValueRange tileSizes, ArrayRef sizeBounds); +/// Add the tile loop induction variables `ivs` to the IndexOp results found in +/// the body of the `tiledOp` to account for the tile offset. +void addTileLoopIvsToIndexOpResults(OpBuilder &b, LinalgOp tiledOp, + ArrayRef ivs); + using FusableOpDependencesTy = llvm::MapVector< Operation *, SmallVector>; diff --git a/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp @@ -222,26 +222,12 @@ } Operation *clonedOp = producer.clone(b, loc, resultTypes, clonedShapes); - // When the producer has index semantics, we have to transform the indices of - // the producer according to the tiling of the consumer, i.e. offset them by - // the values computed in `loopRanges`. - if (producer.hasIndexSemantics()) { - assert(clonedOp->getNumRegions() == 1 && - clonedOp->getRegion(0).getBlocks().size() == 1 && - "expected producer to have one block."); - // Shift all indices by the tile offset. - Block &block = clonedOp->getRegion(0).front(); - for (IndexOp indexOp : block.getOps()) { - OpBuilder::InsertionGuard g(b); - b.setInsertionPointAfter(indexOp); - AffineExpr index, offset; - bindDims(b.getContext(), index, offset); - AffineApplyOp applyOp = b.create( - indexOp.getLoc(), index + offset, - ValueRange{indexOp.getResult(), loopRanges[indexOp.dim()].offset}); - indexOp.getResult().replaceAllUsesExcept(applyOp, applyOp); - } - } + + // Shift all IndexOp results by the tile offset. + SmallVector allIvs; + transform(loopRanges, std::back_inserter(allIvs), + [](Range range) { return range.offset; }); + addTileLoopIvsToIndexOpResults(b, clonedOp, allIvs); return clonedOp; } diff --git a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp @@ -129,28 +129,14 @@ static void transformIndexOps(OpBuilder &b, LinalgOp op, SmallVectorImpl &ivs, const LoopIndexToRangeIndexMap &loopIndexToRangeIndex) { - // Skip operations that have no region attached. - if (op->getNumRegions() == 0) - return; - assert(op->getNumRegions() == 1 && op->getRegion(0).getBlocks().size() == 1 && - "expected linalg operation to have one block."); - Block &block = op->getRegion(0).front(); - - for (IndexOp indexOp : block.getOps()) { - auto rangeIndex = loopIndexToRangeIndex.find(indexOp.dim()); + SmallVector allIvs(op.getNumLoops(), nullptr); + for (auto &en : enumerate(allIvs)) { + auto rangeIndex = loopIndexToRangeIndex.find(en.index()); if (rangeIndex == loopIndexToRangeIndex.end()) continue; - // Offset the index by the value of the corresponding induction variable and - // replace all uses of the previous value. - OpBuilder::InsertionGuard g(b); - b.setInsertionPointAfter(indexOp); - AffineExpr index, iv; - bindDims(b.getContext(), index, iv); - AffineApplyOp applyOp = b.create( - indexOp.getLoc(), index + iv, - ValueRange{indexOp.getResult(), ivs[rangeIndex->second]}); - indexOp.getResult().replaceAllUsesExcept(applyOp, applyOp); + en.value() = ivs[rangeIndex->second]; } + addTileLoopIvsToIndexOpResults(b, op, allIvs); } // Insert a tile `source` into the destination tensor `dest`. The position at diff --git a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp --- a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp +++ b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp @@ -731,5 +731,28 @@ return tiledShapes; } +void addTileLoopIvsToIndexOpResults(OpBuilder &b, LinalgOp tiledOp, + ArrayRef ivs) { + if (tiledOp.hasIndexSemantics()) { + assert(tiledOp->getNumRegions() == 1 && + tiledOp->getRegion(0).getBlocks().size() == 1 && + "expect producer to have one block."); + // Shift all IndexOp results by the tile offset. + Block &block = tiledOp->getRegion(0).front(); + for (IndexOp indexOp : block.getOps()) { + if (ivs[indexOp.dim()] == nullptr) + continue; + OpBuilder::InsertionGuard guard(b); + b.setInsertionPointAfter(indexOp); + AffineExpr index, offset; + bindDims(b.getContext(), index, offset); + AffineApplyOp applyOp = makeComposedAffineApply( + b, indexOp.getLoc(), index + offset, + ValueRange{indexOp.getResult(), ivs[indexOp.dim()]}); + indexOp.getResult().replaceAllUsesExcept(applyOp, applyOp); + } + } +} + } // namespace linalg } // namespace mlir