Skip to content

Commit 7ebfbf9

Browse files
banach-spacebanach-space
authored
[mlir][tensor] Update GeneralizeOuterUnitDimsPackOpPattern (#115312)
Avoid generating spurious tensor.extract_slice, follow-on for #114315. This is best to demonstrate with an example. Here's input for `GeneralizeOuterUnitDimsPackOpPattern`: ```mlir %pack = tensor.pack %input padding_value(%pad : f32) inner_dims_pos = [1, 0] inner_tiles = [2, %tile_dim_1] into %output : tensor<5x1xf32> -> tensor<1x1x2x?xf32> ``` Output _before_: ```mlir %padded = tensor.pad %arg0 low[0, 0] high[%0, 1] { ^bb0(%arg4: index, %arg5: index): tensor.yield %arg2 : f32 } : tensor<5x1xf32> to tensor<?x2xf32> // NOTE: skipped in the output _after_ %extracted_slice = tensor.extract_slice %padded[0, 0] [%arg3, 2] [1, 1] : tensor<?x2xf32> to tensor<?x2xf32> %empty = tensor.empty(%arg3) : tensor<2x?xf32> %transposed = linalg.transpose ins(%extracted_slice : tensor<?x2xf32>) outs(%empty : tensor<2x?xf32>) permutation = [1, 0] %inserted_slice = tensor.insert_slice %transposed= into %arg1[0, 0, 0, 0] [1, 1, 2, %arg3] [1, 1, 1, 1] : tensor<2x?xf32> into tensor<1x1x2x?xf32> ``` Output _after_: ```mlir %padded = tensor.pad %arg0 low[0, 0] high[%0, 1] { ^bb0(%arg4: index, %arg5: index): tensor.yield %arg2 : f32 } : tensor<5x1xf32> to tensor<?x2xf32> %empty = tensor.empty(%arg3) : tensor<2x?xf32> %transposed = linalg.transpose ins(%padded : tensor<?x2xf32>) outs(%empty : tensor<2x?xf32>) permutation = [1, 0] %inserted_slice = tensor.insert_slice %transposed into %arg1[0, 0, 0, 0] [1, 1, 2, %arg3] [1, 1, 1, 1] : tensor<2x?xf32> into tensor<1x1x2x?xf32> ``` This PR also adds a check to verify that only the last N trailing dimensions are tiled (for some value of N). Based on the PR discussion, this restriction seems reasonable - especially as there are no in-tree tests requiring otherwise. For now, it also simplifies the computation of permutations for linalg.transpose. This restriction can be relaxed in the future if needed.
1 parent e458434 commit 7ebfbf9

File tree

4 files changed

+65
-78
lines changed

4 files changed

+65
-78
lines changed

‎mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1516,7 +1516,7 @@ struct GeneralizePadOpPattern : public OpRewritePattern<tensor::PadOp> {
15161516
};
15171517

15181518
/// Rewrites a tensor::PackOp into a sequence of:
1519-
/// * tensor::PadOp + linalg::TransposeOp + tensor::ExtractSliceOp +
1519+
/// * tensor::PadOp + linalg::TransposeOp +
15201520
/// tensor::EmptyOp + tensor::InsertSliceOp ops.
15211521
///
15221522
/// Required that all the outer dims of the input tensor::PackOp are 1.
@@ -1537,10 +1537,6 @@ struct GeneralizePadOpPattern : public OpRewritePattern<tensor::PadOp> {
15371537
/// ^bb0(...):
15381538
/// tensor.yield %arg2 : f32
15391539
/// } : tensor<5x1xf32> to tensor<?x2xf32>
1540-
/// // ExtractSliceOp
1541-
/// %extracted_slice = tensor.extract_slice %padded[0, 0] [%tile_dim_1, 2] [1,
1542-
/// 1]
1543-
/// : tensor<?x2xf32> to tensor<?x2xf32>
15441540
/// // EmptyOp + TransposeOp
15451541
/// %empty = tensor.empty(%arg3) : tensor<2x?xf32>
15461542
/// %transposed = linalg.transpose

‎mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

Lines changed: 40 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -1153,71 +1153,66 @@ LogicalResult GeneralizeOuterUnitDimsPackOpPattern::matchAndRewrite(
11531153
Location loc = packOp.getLoc();
11541154

11551155
Value input = getPackOpSourceOrPaddedSource(rewriter, packOp);
1156-
auto inputShape = packOp.getSourceType().getShape();
11571156
DenseMap<int64_t, OpFoldResult> dimAndTileMapping =
11581157
packOp.getDimAndTileMapping();
11591158
int64_t srcRank = packOp.getSourceRank();
1160-
11611159
int64_t destRank = packOp.getDestRank();
1162-
size_t numTiles = destRank - srcRank;
1163-
1164-
// 1. Use rank-reduced tensor.extract_slice op to extract the tile:
1165-
// %extracted_tile = tensor.extract_slice(%pack_op_input)
1166-
SmallVector<OpFoldResult> readOffsets(srcRank, zeroIdxAttr);
1167-
SmallVector<OpFoldResult> readStrides(srcRank, oneIdxAttr);
1160+
int64_t numTiles = destRank - srcRank;
11681161

1169-
// The sizes attribute for ExtractSliceOp. The leading sizes are set to 1 as
1170-
// all outer dims are 1.
1171-
SmallVector<OpFoldResult> extractSliceSizes(srcRank - numTiles, oneIdxAttr);
1172-
// The shape of the output for ExtractSliceOp. All leading unit dims are
1173-
// effectively rank-reduced, hence skipped.
1174-
SmallVector<int64_t> outputShapeForExtractSlice;
1162+
if (!llvm::all_of(packOp.getInnerDimsPos(),
1163+
[&srcRank, &numTiles](int64_t dimPos) {
1164+
return dimPos >= (srcRank - numTiles - 1);
1165+
}))
1166+
return rewriter.notifyMatchFailure(
1167+
packOp, "Attempting to tile non-trailing source dims!");
11751168

1176-
// Extract the trailing sizes and shape dims for ExtractSliceOp. These should
1177-
// be equal to the inner tile sizes.
1169+
// 1. Extract the inner tile sizes.
1170+
// Where possible, values are replaced with constant attributes (to match the
1171+
// behaviour of `getPackOpSourceOrPaddedSource`).
1172+
SmallVector<OpFoldResult> tileSizes;
11781173
for (auto i : llvm::seq<unsigned>(0, srcRank)) {
11791174
if (dimAndTileMapping.count(i)) {
1180-
auto [tileSize, tileSizeOfr] =
1175+
// Rather than taking the tile size as is, extact the actual constant
1176+
// value Attribute where possible, e.g.:
1177+
// [Value: %tile_size = arith.constant 8 : index] --> [Attribute: 8]
1178+
auto [_, tileSize] =
11811179
getSimplifiedOfrAndStaticSizePair(dimAndTileMapping[i], rewriter);
1182-
extractSliceSizes.push_back(tileSizeOfr);
1183-
outputShapeForExtractSlice.push_back(tileSize);
1180+
tileSizes.push_back(tileSize);
11841181
}
11851182
}
11861183

1187-
Type elemType = packOp.getSourceType().getElementType();
1188-
auto readType = RankedTensorType::get(outputShapeForExtractSlice, elemType);
1189-
1190-
Value tile = rewriter.create<tensor::ExtractSliceOp>(
1191-
loc, readType, input, readOffsets, extractSliceSizes, readStrides);
1192-
1193-
// 2. Transpose the tile to match the inner tile order:
1184+
// 2. Transpose the input to match the inner tile order:
11941185
// %init = tensor.empty()
1195-
// %transposed_tile = linalg.transpose ins(%extracted_tile), outs(%init)
1196-
// NOTE: Outer dims are 1 and hence effectively ignored.
1197-
SmallVector<int64_t> perm = getPackUnpackRankReducedPerm(
1198-
inputShape, packOp.getInnerDimsPos(), packOp.getOuterDimsPerm());
1186+
// %transposed_tile = linalg.transpose ins(%source_or_padded_source),
1187+
// outs(%init)
1188+
// Two assumptions are made:
1189+
// 1. All outer dims are 1 - the corresponding transposition doesn't matter.
1190+
// 2. Inner dims position correspond to the trailing `numTiles` dims.
1191+
SmallVector<int64_t> tilesPermNormalized =
1192+
getPackUnpackNormalizedPerm(srcRank, packOp.getInnerDimsPos());
1193+
SmallVector<int64_t> srcPermForTranspose;
1194+
for (int64_t i = 0; i < (srcRank - numTiles); i++)
1195+
srcPermForTranspose.push_back(i);
1196+
1197+
srcPermForTranspose.append(SmallVector<int64_t>(packOp.getInnerDimsPos()));
11991198

12001199
LLVM_DEBUG(DBGS() << "Pack permutation: " << packOp << "\n";
1201-
llvm::interleaveComma(perm, DBGS() << "perm: "); DBGSNL(););
1200+
llvm::interleaveComma(srcPermForTranspose, DBGS() << "perm: ");
1201+
DBGSNL(););
12021202

12031203
// 2.1 Create tensor.empty (init value for TransposeOp)
1204-
SmallVector<OpFoldResult> transShapeForEmptyOp;
1205-
1206-
// Acquire tensor shape required to create EmptyOp. This will match the inner
1207-
// tile sizes.
1208-
size_t idx = numTiles;
1209-
while (idx != 0) {
1210-
transShapeForEmptyOp.push_back(extractSliceSizes[srcRank - idx]);
1211-
idx--;
1212-
}
1204+
SmallVector<OpFoldResult> transShapeForEmptyOp(srcRank - numTiles,
1205+
oneIdxAttr);
1206+
transShapeForEmptyOp.append(tileSizes);
12131207

1214-
applyPermutationToVector<OpFoldResult>(transShapeForEmptyOp, perm);
1215-
Value empty =
1216-
rewriter.create<tensor::EmptyOp>(loc, transShapeForEmptyOp, elemType);
1208+
applyPermutationToVector<OpFoldResult>(transShapeForEmptyOp,
1209+
srcPermForTranspose);
1210+
Value empty = rewriter.create<tensor::EmptyOp>(
1211+
loc, transShapeForEmptyOp, packOp.getSourceType().getElementType());
12171212

12181213
// 2.2 Create linalg.transpose
1219-
auto transposedOp =
1220-
rewriter.create<linalg::TransposeOp>(loc, tile, empty, perm);
1214+
auto transposedOp = rewriter.create<linalg::TransposeOp>(loc, input, empty,
1215+
srcPermForTranspose);
12211216

12221217
// 3. Insert the inner tile to the destination:
12231218
// %inserted_tile = tensor.insert_slice(%transposed_tile)

‎mlir/test/Dialect/Linalg/generalize-tensor-pack-tile.mlir

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -9,19 +9,17 @@ func.func @KCRS_to_KCRSsr(%arg0: tensor<1x1x128x64xf32>, %arg1: tensor<1x1x4x8x8
99
// CHECK: func.func @KCRS_to_KCRSsr
1010
// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]
1111
// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]
12-
// CHECK: %{{.+}} = scf.for %[[R:[a-zA-Z0-9]+]] =
13-
// CHECK: %{{.+}} = scf.for %[[S:[a-zA-Z0-9]+]] =
12+
// CHECK: scf.for %[[R:[a-zA-Z0-9]+]] =
13+
// CHECK: scf.for %[[S:[a-zA-Z0-9]+]] {{.*}} iter_args(%[[ITER_SLICE:.*]] =
1414
// CHECK: %[[IN_R:.+]] = affine.apply #[[MAP0]](%[[R]])
1515
// CHECK: %[[IN_S:.+]] = affine.apply #[[MAP2]](%[[S]])
1616
// CHECK: %[[SRC_SLICE:.+]] = tensor.extract_slice %[[SRC]]
1717
// CHECK-SAME: [0, 0, %[[IN_R]], %[[IN_S]]] [1, 1, 32, 8] [1, 1, 1, 1]
18-
// CHECK: %[[TILE:.+]] = tensor.extract_slice %[[SRC_SLICE]]
19-
// CHECK-SAME: [0, 0, 0, 0] [1, 1, 32, 8] [1, 1, 1, 1] : tensor<1x1x32x8xf32> to tensor<32x8xf32>
20-
// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<8x32xf32>
21-
// CHECK: %[[TRANSP:.+]] = linalg.transpose
22-
// CHECK-SAME: ins(%[[TILE]]
23-
// CHECK-SAME: outs(%[[EMPTY]]
24-
// CHECK-SAME: permutation = [1, 0]
18+
// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<1x1x8x32xf32>
19+
// CHECK: %[[TRANSP:.*]] = linalg.transpose
20+
// CHECK-SAME: ins(%[[SRC_SLICE]] : tensor<1x1x32x8xf32>)
21+
// CHECK-SAME: outs(%[[EMPTY]] : tensor<1x1x8x32xf32>)
22+
// CHECK-SAME: permutation = [0, 1, 3, 2]
2523
// CHECK: %{{.+}} = tensor.insert_slice %[[TRANSP]] into %{{.+}}
2624

2725
module attributes {transform.with_named_sequence} {

‎mlir/test/Dialect/Linalg/generalize-tensor-pack.mlir

Lines changed: 17 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -63,8 +63,7 @@ func.func @simple_pad_and_pack_dynamic_tile(%input: tensor<5x1xf32>, %output: te
6363
// CHECK: %[[PAD:.*]] = tensor.pad %[[SRC]] low[0, 0] high{{\[}}%[[PAD_HIGH]], 1] {
6464
// CHECK: tensor.yield %[[PAD_VAL]] : f32
6565
// CHECK-NOT: linalg.transpose
66-
// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[PAD:.*]][0, 0] {{\[}}%[[TILE_DIM_0]], 2] [1, 1] : tensor<?x2xf32> to tensor<?x2xf32>
67-
// CHECK: %[[RES:.*]] = tensor.insert_slice %[[SLICE]] into %[[DEST]][0, 0, 0, 0] [1, 1, %[[TILE_DIM_0]], 2] [1, 1, 1, 1] : tensor<?x2xf32> into tensor<1x1x?x2xf32>
66+
// CHECK: %[[RES:.*]] = tensor.insert_slice %[[PAD:.*]] into %[[DEST]][0, 0, 0, 0] [1, 1, %[[TILE_DIM_0]], 2] [1, 1, 1, 1] : tensor<?x2xf32> into tensor<1x1x?x2xf32>
6867
// CHECK: return %[[RES]] : tensor<1x1x?x2xf32>
6968

7069
func.func @simple_pad_and_pack_dynamic_tile_cst(%input: tensor<5x1xf32>, %output: tensor<1x1x?x2xf32>, %pad: f32) -> tensor<1x1x?x2xf32> {
@@ -95,10 +94,10 @@ func.func @simple_pad_and_pack_dynamic_tile_transpose(%input: tensor<5x1xf32>, %
9594
// CHECK: %[[PAD:.*]] = tensor.pad %[[SRC]] low[0, 0] high{{\[}}%[[PAD_HIGH]], 1] {
9695
// CHECK: tensor.yield %[[PAD_VAL]] : f32
9796
// CHECK-NEXT: } : tensor<5x1xf32> to tensor<?x2xf32>
98-
// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[PAD:.*]][0, 0] {{\[}}%[[TILE_DIM_1]], 2] [1, 1] : tensor<?x2xf32> to tensor<?x2xf32>
9997
// CHECK: %[[EMPTY:.*]] = tensor.empty(%[[TILE_DIM_1]]) : tensor<2x?xf32>
10098
// CHECK: %[[TR:.*]] = linalg.transpose
101-
// CHECK-SAME: ins(%[[SLICE]] : tensor<?x2xf32>) outs(%[[EMPTY]] : tensor<2x?xf32>)
99+
// CHECK-SAME: ins(%[[PAD:.*]] : tensor<?x2xf32>)
100+
// CHECK-SAME: outs(%[[EMPTY]] : tensor<2x?xf32>)
102101
// CHECK-SAME: permutation = [1, 0]
103102
// CHECK: %[[RES:.*]] = tensor.insert_slice %[[TR]] into %[[DEST]][0, 0, 0, 0] [1, 1, 2, %[[TILE_DIM_1]]] [1, 1, 1, 1] : tensor<2x?xf32> into tensor<1x1x2x?xf32>
104103
// CHECK: return %[[RES]] : tensor<1x1x2x?xf32>
@@ -128,10 +127,10 @@ func.func @simple_pad_and_pack_scalable_tile(%input: tensor<5x1xf32>, %output: t
128127
// CHECK: %[[PAD:.+]] = tensor.pad %[[SRC]] low[0, 0] high{{\[}}%[[PAD_HIGH]], 1] {
129128
// CHECK: tensor.yield %[[PAD_VAL]] : f32
130129
// CHECK-NOT: linalg.transpose
131-
// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[PAD:.+]][0, 0] {{\[}}%[[C8_VS]], 2] [1, 1] : tensor<?x2xf32> to tensor<?x2xf32>
132-
// CHECK: %[[RES:.+]] = tensor.insert_slice %[[SLICE]] into %[[DEST]][0, 0, 0, 0] [1, 1, %[[C8_VS]], 2] [1, 1, 1, 1] : tensor<?x2xf32> into tensor<1x1x?x2xf32>
130+
// CHECK: %[[RES:.*]] = tensor.insert_slice %[[PAD:.*]] into %[[DEST]][0, 0, 0, 0] [1, 1, %[[C8_VS]], 2] [1, 1, 1, 1] : tensor<?x2xf32> into tensor<1x1x?x2xf32>
133131
// CHECK: return %[[RES]] : tensor<1x1x?x2xf32>
134132

133+
135134
/// Same as example above, but with both tile sizes dynamic.
136135

137136
func.func @simple_pad_and_pack_dynamic_tiles(%input: tensor<5x1xf32>, %output: tensor<1x1x?x?xf32>, %pad: f32, %tile_dim_0: index, %tile_dim_1: index) -> tensor<1x1x?x?xf32> {
@@ -149,8 +148,7 @@ func.func @simple_pad_and_pack_dynamic_tiles(%input: tensor<5x1xf32>, %output: t
149148
// CHECK: %[[PAD:.*]] = tensor.pad %[[SRC]] low[0, 0] high{{\[}}%[[PAD_HIGH_1]], %[[PAD_HIGH_2]]] {
150149
// CHECK: tensor.yield %[[PAD_VAL]] : f32
151150
// CHECK-NOT: linalg.transpose
152-
// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[PAD:.*]][0, 0] {{\[}}%[[TILE_DIM_0]], %[[TILE_DIM_1]]] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
153-
// CHECK: %[[RES:.*]] = tensor.insert_slice %[[SLICE]] into %[[DEST]][0, 0, 0, 0] [1, 1, %[[TILE_DIM_0]], %[[TILE_DIM_1]]] [1, 1, 1, 1] : tensor<?x?xf32> into tensor<1x1x?x?xf32>
151+
// CHECK: %[[RES:.*]] = tensor.insert_slice %[[PAD]] into %[[DEST]][0, 0, 0, 0] [1, 1, %[[TILE_DIM_0]], %[[TILE_DIM_1]]] [1, 1, 1, 1] : tensor<?x?xf32> into tensor<1x1x?x?xf32>
154152
// CHECK: return %[[RES]] : tensor<1x1x?x?xf32>
155153

156154
// -----
@@ -170,12 +168,13 @@ func.func @simple_pad_and_pack_dynamic_tile_not_all_dims_tiled(%input: tensor<1x
170168
// CHECK: ^bb0(%[[VAL_6:.*]]: index, %[[VAL_7:.*]]: index, %[[VAL_8:.*]]: index, %[[VAL_9:.*]]: index):
171169
// CHECK: tensor.yield %[[VAL_2]] : f32
172170
// CHECK: } : tensor<1x1x5x1xf32> to tensor<1x1x?x2xf32>
173-
// CHECK: %[[VAL_10:.*]] = tensor.extract_slice %[[VAL_11:.*]][0, 0, 0, 0] [1, 1, %[[VAL_3]], 2] [1, 1, 1, 1] : tensor<1x1x?x2xf32> to tensor<?x2xf32>
174-
// CHECK: %[[VAL_12:.*]] = tensor.empty(%[[VAL_3]]) : tensor<2x?xf32>
175-
// CHECK: %[[VAL_13:.*]] = linalg.transpose ins(%[[VAL_10]] : tensor<?x2xf32>) outs(%[[VAL_12]] : tensor<2x?xf32>) permutation = [1, 0]
176-
// CHECK: %[[VAL_14:.*]] = tensor.insert_slice %[[VAL_13]] into %[[VAL_1]][0, 0, 0, 0, 0, 0] [1, 1, 1, 1, 2, %[[VAL_3]]] [1, 1, 1, 1, 1, 1] : tensor<2x?xf32> into tensor<1x1x1x1x2x?xf32>
177-
// CHECK: return %[[VAL_14]] : tensor<1x1x1x1x2x?xf32>
178-
// CHECK: }
171+
// CHECK: %[[VAL_10:.*]] = tensor.empty(%[[VAL_3]]) : tensor<1x1x2x?xf32>
172+
// CHECK: %[[VAL_11:.*]] = linalg.transpose
173+
// CHECK-SAME: ins(%[[VAL_12:.*]] : tensor<1x1x?x2xf32>)
174+
// CHECK-SAME: outs(%[[VAL_10]] : tensor<1x1x2x?xf32>)
175+
// CHECK-SAME: permutation = [0, 1, 3, 2]
176+
// CHECK: %[[VAL_13:.*]] = tensor.insert_slice %[[VAL_11]] into %[[VAL_1]][0, 0, 0, 0, 0, 0] [1, 1, 1, 1, 2, %[[VAL_3]]] [1, 1, 1, 1, 1, 1] : tensor<1x1x2x?xf32> into tensor<1x1x1x1x2x?xf32>
177+
// CHECK: return %[[VAL_13]] : tensor<1x1x1x1x2x?xf32>
179178

180179
// -----
181180

@@ -218,12 +217,11 @@ func.func @simple_KCRS_to_KRSCsr(%arg0: tensor<1x1x32x8xf32>, %arg1: tensor<1x1x
218217
// CHECK-LABEL: func.func @simple_KCRS_to_KRSCsr
219218
// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]
220219
// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]
221-
// CHECK: %[[TILE:.+]] = tensor.extract_slice %[[SRC]][0, 0, 0, 0] [1, 1, 32, 8] [1, 1, 1, 1]
222-
// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<8x32xf32>
220+
// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<1x1x8x32xf32>
223221
// CHECK: %[[TRANSP:.+]] = linalg.transpose
224-
// CHECK-SAME: ins(%[[TILE]] : tensor<32x8xf32>)
225-
// CHECK-SAME: outs(%[[EMPTY]] : tensor<8x32xf32>)
226-
// CHECK-SAME: permutation = [1, 0]
222+
// CHECK-SAME: ins(%[[SRC]] : tensor<1x1x32x8xf32>
223+
// CHECK-SAME: outs(%[[EMPTY]] : tensor<1x1x8x32xf32>)
224+
// CHECK-SAME: permutation = [0, 1, 3, 2]
227225
// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[TRANSP]] into %[[DEST]]
228226
// CHECK-SAME: [0, 0, 0, 0, 0, 0] [1, 1, 1, 1, 8, 32] [1, 1, 1, 1, 1, 1]
229227
// CHECK: return %[[INSERT]]

0 commit comments

Comments
 (0)