[Codegen] Support more operations in tile size analysis#23971
[Codegen] Support more operations in tile size analysis#23971sommerlukas wants to merge 3 commits intoiree-org:mainfrom
Conversation
Signed-off-by: Lukas Sommer <lukas.sommer@amd.com>
Signed-off-by: Lukas Sommer <lukas.sommer@amd.com>
sommerlukas
requested review from
MaheshRavishankar,
Max191,
hanhanW and
qedawkins
as code owners
sommerlukas
changed the title
hanhanW
left a comment
hanhanW
left a comment
There was a problem hiding this comment.
I feel that dynamic cases are not exercised enough, because I'm seeing weird assumption in the implementation. E.g., the inlined comment in mapPackSourceToDest.
Not saying dynamic inner tile sizes, the inner tile sizes are static in the scope. The outer dimension can be dynamic and the op can carry padding value.
There was a problem hiding this comment.
I think we should have tests with dynamic shapes.
| if (!outerDimsPerm.empty()) { | ||
| applyPermutationToVector(result.dims, outerDimsPerm); | ||
| } |
There was a problem hiding this comment.
We need a test for outer_dims_perm. Please use a permutation that is not identical to the inversed permutation. e.g., do not use [1, 0].
| for (int64_t t : staticInnerTiles) { | ||
| result.dims.push_back(t); | ||
| } |
There was a problem hiding this comment.
result.dims.append(staticInnerTiles)?
| for (auto [dimPos, tileSize] : | ||
| llvm::zip_equal(innerDimsPos, staticInnerTiles)) { | ||
| if (result.dims[dimPos] == kUninitialized || | ||
| result.dims[dimPos] == kOverdefined) { | ||
| return {}; | ||
| } | ||
| if (result.dims[dimPos] % tileSize != 0) { | ||
| return {}; | ||
| } | ||
| result.dims[dimPos] /= tileSize; | ||
| } |
There was a problem hiding this comment.
Q1: Don't we need to check if tileSize is static or not?
Q2: don't we preserve kUninitialized and kOverdefined when it happens? Why do you return {}?
There was a problem hiding this comment.
Q1: This is already checked at the call-sites.
Q2: Changed to preserve.
There was a problem hiding this comment.
RE Q1: If we look the code from class's perspective, it is better to document the assumption in method comment, or we should just handle it here.
| /// Append dimensions from `suffix` to produce a higher-rank TileSizes. | ||
| TileSizes extend(ArrayRef<int64_t> suffix) const { | ||
| SmallVector<int64_t> fullDims(dims); | ||
| fullDims.append(suffix.begin(), suffix.end()); | ||
| return TileSizes(fullDims); | ||
| } |
There was a problem hiding this comment.
How about naming it as append which matches the comment and methods from SmallVector?
| linalgOp.emitOpError() | ||
| << "tile size analysis did not determine a valid tile size"; | ||
| auto materialize = | ||
| [&](Operation *op, TileSizes tileSizes) { |
There was a problem hiding this comment.
Do you need to capture anything?
| [&](Operation *op, TileSizes tileSizes) { | |
| [](Operation *op, TileSizes tileSizes) { |
| static TileSizes getIterationSpaceTileSizes(Operation *op, unsigned numLoops, | ||
| ArrayRef<AffineMap> indexingMaps, | ||
| const DataFlowSolver &solver) { |
There was a problem hiding this comment.
I feel that you want to pass in TilingInterface op which has the additional information.
There was a problem hiding this comment.
IIUC, TilingInterface doesn't give us access to the indexing maps, so we won't gain much here.
There was a problem hiding this comment.
Oh, I missed that. I was going to suggest IndexingMapOpInterface, then I found that InnerTileOp haven't implemented the interface, so it is okay for now.
| if (auto linalgOp = dyn_cast<linalg::LinalgOp>(op)) { | ||
| SmallVector<AffineMap> indexingMaps = linalgOp.getIndexingMapsArray(); | ||
| TileSizes tileSizes = getIterationSpaceTileSizes( | ||
| op, linalgOp.getNumLoops(), indexingMaps, solver); | ||
| assert(!tileSizes.isDefined() || | ||
| tileSizes.rank() == linalgOp.getNumLoops()); | ||
| materialize(op, tileSizes); | ||
| return; | ||
| } | ||
| if (!tileSizes.isDefined()) { | ||
| LDBG() << "Analysis did not determine tile size for " << *linalgOp; | ||
|
|
||
| if (auto innerTiledOp = dyn_cast<IREE::Codegen::InnerTiledOp>(op)) { | ||
| SmallVector<AffineMap> indexingMaps = | ||
| innerTiledOp.getIndexingMapsArray(); | ||
| unsigned numLoops = indexingMaps[0].getNumDims(); | ||
| TileSizes tileSizes = | ||
| getIterationSpaceTileSizes(op, numLoops, indexingMaps, solver); | ||
| materialize(op, tileSizes); | ||
| return; | ||
| } |
There was a problem hiding this comment.
Like the other comment, I feel that we should structure these like:
// Special cases like pack/unpack
// TilingInterface op. Or isa<linalg::LinalgOp, InnerTiledOp>
| op->emitOpError() | ||
| << "tile size analysis did not determine a valid tile size"; | ||
| return; |
There was a problem hiding this comment.
I missed this when you landed the previous PR. I was not clear enough. I think the pass should signal failure if it is overdefined. Emitting error messages is bundled with pass failure to me, and I did not explicitly spell it out in my other comment.
There was a problem hiding this comment.
Added proper pass failure.
| if (result.dims[dimPos] % tileSize != 0) { | ||
| return {}; | ||
| } | ||
| result.dims[dimPos] /= tileSize; |
There was a problem hiding this comment.
I thought that it is always ceilDiv in this case, because pack have padding semantics?
There was a problem hiding this comment.
I've switched to using ceil division, but in the same course also deactivated backward propagation for pack with padding.
That propagation would in the padding case would otherwise propagate too large tile sizes (or lead to overdefined).
compiler/src/iree/compiler/Codegen/Common/MaterializeVectorTileSizes.cpp
Show resolved
Hide resolved
Signed-off-by: Lukas Sommer <lukas.sommer@amd.com>
sommerlukas
left a comment
sommerlukas
left a comment
There was a problem hiding this comment.
Thanks for the feedback! I've added tests for dynamic shapes and outer_dims_perm, some more comments inline.
| if (result.dims[dimPos] % tileSize != 0) { | ||
| return {}; | ||
| } | ||
| result.dims[dimPos] /= tileSize; |
There was a problem hiding this comment.
I've switched to using ceil division, but in the same course also deactivated backward propagation for pack with padding.
That propagation would in the padding case would otherwise propagate too large tile sizes (or lead to overdefined).
| return success(); | ||
| } | ||
|
|
||
| // InnerTiledOp: propagate through indexing maps (outer dims only). |
There was a problem hiding this comment.
The comment isn't phrased very well. We don't exclude the inner dims (we append them before propagating, see below), but they don't participate in the propagation through indexing maps (because they don't need to). I've clarified the comment. We also don't materialize the inner dims in the attribute on inner_tiled, because it would duplicate the information.
| // Uninitialized (0) stays 0 after multiply; overdefined is preserved | ||
| // by skipping. Safe because we only multiply here, unlike | ||
| // mapPackSourceToDest which must divide (and can't divide by zero). | ||
| continue; |
| for (auto [dimPos, tileSize] : | ||
| llvm::zip_equal(innerDimsPos, staticInnerTiles)) { | ||
| if (result.dims[dimPos] == kUninitialized || | ||
| result.dims[dimPos] == kOverdefined) { | ||
| return {}; | ||
| } | ||
| if (result.dims[dimPos] % tileSize != 0) { | ||
| return {}; | ||
| } | ||
| result.dims[dimPos] /= tileSize; | ||
| } |
There was a problem hiding this comment.
Q1: This is already checked at the call-sites.
Q2: Changed to preserve.
| /// Append dimensions from `suffix` to produce a higher-rank TileSizes. | ||
| TileSizes extend(ArrayRef<int64_t> suffix) const { | ||
| SmallVector<int64_t> fullDims(dims); | ||
| fullDims.append(suffix.begin(), suffix.end()); | ||
| return TileSizes(fullDims); | ||
| } |
| op->emitOpError() | ||
| << "tile size analysis did not determine a valid tile size"; | ||
| return; |
There was a problem hiding this comment.
Added proper pass failure.
| static TileSizes getIterationSpaceTileSizes(Operation *op, unsigned numLoops, | ||
| ArrayRef<AffineMap> indexingMaps, | ||
| const DataFlowSolver &solver) { |
There was a problem hiding this comment.
IIUC, TilingInterface doesn't give us access to the indexing maps, so we won't gain much here.
| for (auto [dimPos, tileSize] : | ||
| llvm::zip_equal(innerDimsPos, staticInnerTiles)) { | ||
| if (result.dims[dimPos] == kUninitialized || | ||
| result.dims[dimPos] == kOverdefined) { | ||
| return {}; | ||
| } | ||
| if (result.dims[dimPos] % tileSize != 0) { | ||
| return {}; | ||
| } | ||
| result.dims[dimPos] /= tileSize; | ||
| } |
There was a problem hiding this comment.
RE Q1: If we look the code from class's perspective, it is better to document the assumption in method comment, or we should just handle it here.
| return success(); | ||
| } | ||
|
|
||
| // InnerTiledOp: propagate through indexing maps (outer dims only). |
There was a problem hiding this comment.
I've clarified the comment.
I don't see the update on the comment. Do you miss the change?
| if (ShapedType::isDynamicShape(packOp.getStaticInnerTiles())) { | ||
| return success(); | ||
| } | ||
| if (packOp.getPaddingValue()) { |
There was a problem hiding this comment.
Can we have a test for padding value?
| static TileSizes getIterationSpaceTileSizes(Operation *op, unsigned numLoops, | ||
| ArrayRef<AffineMap> indexingMaps, | ||
| const DataFlowSolver &solver) { |
There was a problem hiding this comment.
Oh, I missed that. I was going to suggest IndexingMapOpInterface, then I found that InnerTileOp haven't implemented the interface, so it is okay for now.
3 participants
Add support for
inner_tiled,linalg.packandlinalg.unpackoperations to the vector tile size analysis.This is part of #23415.
Assisted-by: Claude Code