Skip to content

Commit 3f136f7

Browse files
[Tensor] Simplify tenor.pad tiling length calculations. (#119039)
The current calculations calculate ending location of the new length and then subtract the new offset from that location. It is possible to directly calculate new length. Along with requiring less operations (which can matter in dynamic case) this also has the advantage that the values are upper bounded by length rather than source size which is more friendly for range analysis. I believe the change is already being tested by `test/Dialect/Linalg/subtensor-of-padtensor.mlir` and `test/Dialect/Linalg/tile-and-fuse-tensors.mlir` --------- Signed-off-by: Nirvedh <[email protected]>
1 parent 61510b5 commit 3f136f7

File tree

1 file changed

+8
-15
lines changed

1 file changed

+8
-15
lines changed

mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp

Lines changed: 8 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -746,11 +746,6 @@ FailureOr<TilingResult> tensor::bubbleUpPadSlice(OpBuilder &b,
746746
Location loc = padOp->getLoc();
747747
AffineExpr dim0, dim1;
748748
bindDims(b.getContext(), dim0, dim1);
749-
// Add two integers.
750-
auto addMap = AffineMap::get(2, 0, {dim0 + dim1});
751-
auto add = [&](OpFoldResult v1, OpFoldResult v2) {
752-
return affine::makeComposedFoldedAffineApply(b, loc, addMap, {v1, v2});
753-
};
754749
// Subtract two integers.
755750
auto subMap = AffineMap::get(2, 0, {dim0 - dim1});
756751
auto sub = [&](OpFoldResult v1, OpFoldResult v2) {
@@ -825,16 +820,14 @@ FailureOr<TilingResult> tensor::bubbleUpPadSlice(OpBuilder &b,
825820
// The original read could also have stopped in the high padding zone.
826821
// In that case, set the end positition of the read should be the end of
827822
// the source tensor. (Similar to newOffset.)
828-
//
829-
// endLoc = min(max(offset - low + length, 0), srcSize)
830-
//
831-
// The new ExtractSliceOp length is `endLoc - newOffset`.
832-
//
833-
// Optimization: If low = 0, then the formula can be simplified.
834-
OpFoldResult endLoc =
835-
hasLowPad ? min(max(add(sub(offset, low), length), zero), srcSize)
836-
: min(add(offset, length), srcSize);
837-
OpFoldResult newLength = sub(endLoc, newOffset);
823+
// srcSize - newOffset represents how much length we have available
824+
// and length - newLow represents how much length we want at most.
825+
// Note that there are many ways to order this indexing math to compute newLength, but we want to make sure that the final affine.min ops in the sequence are bounding the index to as small a value as possible. If ValueBoundsOpInterface is used, this calcuation will get upper bounds from the affine.min ops, so we want to use the smallest known value to set the bound at the end of the computation sequence. In this case, the index will be upper bounded by length - newLow.
826+
OpFoldResult newLength = min(sub(srcSize, newOffset), sub(length, newLow));
827+
// Optimization: If low = 0, then newLow = 0. then newLength >= 0 assuming
828+
// length >= 0.
829+
if (hasLowPad)
830+
newLength = max(newLength, zero);
838831
newLengths.push_back(newLength);
839832

840833
// Check if newLength is zero. In that case, no SubTensorOp should be

0 commit comments

Comments
 (0)