From 90b7817e03af19a8fdc7f32f82e39d0fbf8a9791 Mon Sep 17 00:00:00 2001 From: Tobias Gysi Date: Fri, 17 Sep 2021 14:54:04 +0000 Subject: [PATCH] [mlir][linalg] Add helper to update IndexOps after tiling (NFC). Add the addTileLoopIvsToIndexOpResults method to shift the IndexOp results after tiling. Reviewed By: nicolasvasilache Differential Revision: https://reviews.llvm.org/D109761 --- mlir/include/mlir/Dialect/Linalg/Utils/Utils.h | 5 +++++ mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp | 26 ++++++-------------------- mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp | 24 +++++------------------- mlir/lib/Dialect/Linalg/Utils/Utils.cpp | 23 +++++++++++++++++++++++ 4 files changed, 39 insertions(+), 39 deletions(-) diff --git a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h index fd0a0be..8776b74 100644 --- a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h +++ b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h @@ -127,6 +127,11 @@ SmallVector makeTiledShapes(OpBuilder &builder, Location loc, 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 index 7155690..adce2c7 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp @@ -222,26 +222,12 @@ static LinalgOp fuse(OpBuilder &b, LinalgOp producer, } 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 index 57e77c0..f3b9516 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp @@ -129,28 +129,14 @@ makeTiledLoopRanges(OpBuilder &b, Location loc, AffineMap map, 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 index 296956d..e0b1a48 100644 --- a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp +++ b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp @@ -731,5 +731,28 @@ SmallVector makeTiledShapes(OpBuilder &b, Location loc, 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 -- 2.7.4