Skip to content

Commit 75f7295

Browse files
qedawkinsqedawkins
authored
[mlir][Tensor] Fix unpack -> transpose folding pattern for padded unpacks (#90678)
Previously if the producer tensor.unpack op had "unpadding" semantics, the folding pattern would construct a destination that does not match with the result type of the transpose. Because both ops are DPS we can just reuse the destination of the transpose. Additionally cleans up a bunch of trailing whitespace in the test file.
1 parent a03eeb0 commit 75f7295

File tree

2 files changed

+88
-64
lines changed

2 files changed

+88
-64
lines changed

‎mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -374,13 +374,10 @@ struct FoldProducerUnPackWithConsumerLinalgTransposeOp
374374
for (auto dim : innerDimsPos)
375375
newInnerDimsPosVec.push_back(transposePermutation[dim]);
376376

377-
Value output = unPackOp.createDestinationTensor(
378-
rewriter, transposeOp.getLoc(), unPackOp.getSource(),
379-
unPackOp.getMixedTiles(), newInnerDimsPosVec, newOuterDimsPermVec);
380-
377+
// Reuse the destination of the transpose op.
381378
rewriter.replaceOpWithNewOp<UnPackOp>(
382-
transposeOp, unPackOp.getSource(), output, newInnerDimsPosVec,
383-
unPackOp.getMixedTiles(), newOuterDimsPermVec);
379+
transposeOp, unPackOp.getSource(), transposeOp.getDpsInits()[0],
380+
newInnerDimsPosVec, unPackOp.getMixedTiles(), newOuterDimsPermVec);
384381

385382
return success();
386383
}

‎mlir/test/Dialect/Tensor/fold-into-pack-and-unpack.mlir

Lines changed: 85 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ func.func @tensor_pack_linalg_transpose_fold(%arg0: tensor<56x57x1x64xf32>) -> t
137137
// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x57x56x2x32xf32>
138138
// CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]]
139139
// CHECK-SAME: outer_dims_perm = [2, 1, 0, 3]
140-
// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]
140+
// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]
141141
// CHECK-SAME: into %[[INIT]]
142142
// CHECK: return %[[PACK]]
143143

@@ -163,7 +163,7 @@ func.func @tensor_pack_linalg_transpose_fold_with_padding(%arg0: tensor<56x57x1x
163163
// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x57x56x2x32xf32>
164164
// CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]] padding_value(%[[PADDING]] : f32)
165165
// CHECK-SAME: outer_dims_perm = [2, 1, 0, 3]
166-
// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]
166+
// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]
167167
// CHECK-SAME: into %[[INIT]]
168168
// CHECK: return %[[PACK]]
169169

@@ -188,7 +188,7 @@ func.func @tensor_pack_linalg_transpose_fold_no_outer_dims_perm(%arg0: tensor<56
188188
// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x2x56x57x32xf32>
189189
// CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]]
190190
// CHECK-SAME: outer_dims_perm = [2, 3, 0, 1]
191-
// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]
191+
// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]
192192
// CHECK-SAME: into %[[INIT]]
193193
// CHECK: return %[[PACK]]
194194

@@ -214,7 +214,7 @@ func.func @tensor_pack_linalg_transpose_fold_tile_dims_transpose(%arg0: tensor<5
214214
// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<12x56x4x9x32x8x2xf32>
215215
// CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]]
216216
// CHECK-SAME: outer_dims_perm = [2, 0, 3, 1]
217-
// CHECK-SAME: inner_dims_pos = [3, 1, 2] inner_tiles = [32, 8, 2]
217+
// CHECK-SAME: inner_dims_pos = [3, 1, 2] inner_tiles = [32, 8, 2]
218218
// CHECK-SAME: into %[[INIT]]
219219
// CHECK: return %[[PACK]]
220220

@@ -256,7 +256,7 @@ func.func @tensor_pack_linalg_transpose_fold_dynamic_outer_dims(%arg0: tensor<56
256256
outs(%1 : tensor<1x57x56x2x32xf32>)
257257
permutation = [2, 3, 0, 1, 4]
258258

259-
%return_value = tensor.cast %transposed : tensor<1x57x56x2x32xf32> to tensor<?x?x56x2x32xf32>
259+
%return_value = tensor.cast %transposed : tensor<1x57x56x2x32xf32> to tensor<?x?x56x2x32xf32>
260260
return %return_value : tensor<?x?x56x2x32xf32>
261261
}
262262
// CHECK: func @tensor_pack_linalg_transpose_fold_dynamic_outer_dims(
@@ -268,7 +268,7 @@ func.func @tensor_pack_linalg_transpose_fold_dynamic_outer_dims(%arg0: tensor<56
268268
// CHECK: %[[INIT:.+]] = tensor.empty(%[[dim_0]], %[[dim]]) : tensor<?x?x56x2x32xf32>
269269
// CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]]
270270
// CHECK-SAME: outer_dims_perm = [2, 1, 0, 3]
271-
// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]
271+
// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]
272272
// CHECK-SAME: into %[[INIT]]
273273
// CHECK: return %[[PACK]]
274274

@@ -286,14 +286,14 @@ func.func @tensor_pack_linalg_transpose_fold_dynamic_outer_and_tile_dims(%arg0:
286286
ins(%pack : tensor<56x9x12x4x8x2x32xf32>)
287287
outs(%1 : tensor<12x4x56x9x32x8x2xf32>)
288288
permutation = [2, 3, 0, 1, 6, 4, 5]
289-
290-
%return_value = tensor.cast %transposed : tensor<12x4x56x9x32x8x2xf32> to tensor<?x?x56x9x32x8x2xf32>
289+
290+
%return_value = tensor.cast %transposed : tensor<12x4x56x9x32x8x2xf32> to tensor<?x?x56x9x32x8x2xf32>
291291
return %return_value : tensor<?x?x56x9x32x8x2xf32>
292292
}
293293
// CHECK-DAG: #[[$MAP0:.+]] = affine_map<()[s0] -> (s0 ceildiv 8)>
294294
// CHECK-DAG: #[[$MAP1:.+]] = affine_map<()[s0] -> (s0 ceildiv 2)>
295295
// CHECK-LABEL: func.func @tensor_pack_linalg_transpose_fold_dynamic_outer_and_tile_dims(
296-
// CHECK-SAME: %[[ARG0:.+]]: tensor<56x?x?x128xf32>)
296+
// CHECK-SAME: %[[ARG0:.+]]: tensor<56x?x?x128xf32>)
297297
// CHECK-DAG: %[[c1:.+]] = arith.constant 1 : index
298298
// CHECK-DAG: %[[c2:.+]] = arith.constant 2 : index
299299
// CHECK: %[[dim:.+]] = tensor.dim %[[ARG0]], %[[c1]] : tensor<56x?x?x128xf32>
@@ -328,7 +328,7 @@ func.func @tensor_pack_linalg_transpose_fold_dynamic_outer_dims_tile_dims_tile_s
328328
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?xf32>,
329329
// CHECK-SAME: %[[PACK_DEST:.+]]: tensor<?x?x?x?x?x?x?xf32>, %[[TRANSPOSE_DEST:.+]]: tensor<?x?x?x?x?x?x?xf32>,
330330
// CHECK-SAME: %[[ARG1:.+]]: index, %[[ARG2:.+]]: index,
331-
// CHECK-SAME: %[[ARG3:.+]]: index)
331+
// CHECK-SAME: %[[ARG3:.+]]: index)
332332
// CHECK-DAG: %[[c0:.+]] = arith.constant 0 : index
333333
// CHECK-DAG: %[[c1:.+]] = arith.constant 1 : index
334334
// CHECK-DAG: %[[c2:.+]] = arith.constant 2 : index
@@ -367,7 +367,7 @@ func.func @linalg_transpose_tensor_pack_fold(%arg0: tensor<56x57x1x64xf32>) -> t
367367
// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x57x56x2x32xf32>
368368
// CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]]
369369
// CHECK-SAME: outer_dims_perm = [2, 1, 0, 3]
370-
// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]
370+
// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]
371371
// CHECK-SAME: into %[[INIT]]
372372
// CHECK: return %[[PACK]]
373373

@@ -379,7 +379,7 @@ func.func @linalg_transpose_tensor_pack_fold_with_padding(%arg0: tensor<56x57x1x
379379
ins(%arg0 : tensor<56x57x1x55xf32>)
380380
outs(%0 : tensor<1x56x57x55xf32>)
381381
permutation = [2, 0, 1, 3]
382-
382+
383383
%1 = tensor.empty() : tensor<1x57x56x2x32xf32>
384384
%pack = tensor.pack %transpose padding_value(%padding : f32)
385385
outer_dims_perm = [0, 2, 1, 3]
@@ -393,7 +393,7 @@ func.func @linalg_transpose_tensor_pack_fold_with_padding(%arg0: tensor<56x57x1x
393393
// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x57x56x2x32xf32>
394394
// CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]] padding_value(%[[PADDING]] : f32)
395395
// CHECK-SAME: outer_dims_perm = [2, 1, 0, 3]
396-
// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]
396+
// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]
397397
// CHECK-SAME: into %[[INIT]]
398398
// CHECK: return %[[PACK]]
399399

@@ -405,7 +405,7 @@ func.func @linalg_transpose_tensor_pack_fold_no_outer_dims_perm(%arg0: tensor<56
405405
ins(%arg0 : tensor<56x57x1x64xf32>)
406406
outs(%0 : tensor<1x56x57x64xf32>)
407407
permutation = [2, 0, 1, 3]
408-
408+
409409
%1 = tensor.empty() : tensor<1x56x57x2x32xf32>
410410
%pack = tensor.pack %transposed
411411
inner_dims_pos = [3]
@@ -418,7 +418,7 @@ func.func @linalg_transpose_tensor_pack_fold_no_outer_dims_perm(%arg0: tensor<56
418418
// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x56x57x2x32xf32>
419419
// CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]]
420420
// CHECK-SAME: outer_dims_perm = [2, 0, 1, 3]
421-
// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]
421+
// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]
422422
// CHECK-SAME: into %[[INIT]]
423423
// CHECK: return %[[PACK]]
424424

@@ -429,7 +429,7 @@ func.func @linalg_transpose_tensor_pack_fold_complex_inner_dims_change(%arg0: te
429429
ins(%arg0 : tensor<25x30x35x40xf32>)
430430
outs(%transpose_dest : tensor<35x40x25x30xf32>)
431431
permutation = [2, 3, 0, 1]
432-
432+
433433
%pack = tensor.pack %transposed
434434
outer_dims_perm = [3, 0, 2, 1]
435435
inner_dims_pos = [1, 3, 2]
@@ -438,15 +438,15 @@ func.func @linalg_transpose_tensor_pack_fold_complex_inner_dims_change(%arg0: te
438438
return %pack : tensor<3x35x5x8x5x10x5xf32>
439439
}
440440
//CHECK-LABEL: func.func @linalg_transpose_tensor_pack_fold_complex_inner_dims_change(
441-
// CHECK-SAME: %[[ARG0:.+]]: tensor<25x30x35x40xf32>,
442-
// CHECK-SAME: %[[ARG1:.+]]: tensor<35x40x25x30xf32>,
441+
// CHECK-SAME: %[[ARG0:.+]]: tensor<25x30x35x40xf32>,
442+
// CHECK-SAME: %[[ARG1:.+]]: tensor<35x40x25x30xf32>,
443443
// CHECK-SAME: %[[ARG2:.+]]: tensor<3x35x5x8x5x10x5xf32>) -> tensor<3x35x5x8x5x10x5xf32> {
444444
// CHECK: %[[VAL0:.+]] = tensor.empty() : tensor<3x35x5x8x5x10x5xf32>
445-
// CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]]
446-
// CHECK-SAME: outer_dims_perm = [1, 2, 0, 3]
447-
// CHECK-SAME: inner_dims_pos = [3, 1, 0]
448-
// CHECK-SAME: inner_tiles = [5, 10, 5]
449-
// CHECK-SAME: into %[[VAL0]]
445+
// CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]]
446+
// CHECK-SAME: outer_dims_perm = [1, 2, 0, 3]
447+
// CHECK-SAME: inner_dims_pos = [3, 1, 0]
448+
// CHECK-SAME: inner_tiles = [5, 10, 5]
449+
// CHECK-SAME: into %[[VAL0]]
450450
// CHECK: return %[[PACK]]
451451

452452
// -----
@@ -456,7 +456,7 @@ func.func @linalg_transpose_tensor_pack_fold_dynamic_outer_dims_tile_dims_tile_s
456456
ins(%arg0 : tensor<?x?x?x?xf32>)
457457
outs(%transpose_dest : tensor<?x?x?x?xf32>)
458458
permutation = [2, 3, 0, 1]
459-
459+
460460
%pack = tensor.pack %transposed
461461
outer_dims_perm = [3, 0, 2, 1]
462462
inner_dims_pos = [1, 3, 2]
@@ -466,7 +466,7 @@ func.func @linalg_transpose_tensor_pack_fold_dynamic_outer_dims_tile_dims_tile_s
466466
}
467467
// CHECK: #[[$MAP:.+]] = affine_map<()[s0, s1] -> (s0 ceildiv s1)>
468468
//CHECK-LABEL: func.func @linalg_transpose_tensor_pack_fold_dynamic_outer_dims_tile_dims_tile_sizes(
469-
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?xf32>, %[[ARG1:.+]]: tensor<?x?x?x?xf32>,
469+
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?xf32>, %[[ARG1:.+]]: tensor<?x?x?x?xf32>,
470470
// CHECK-SAME: %[[ARG2:.+]]: tensor<?x?x?x?x?x?x?xf32>, %[[ARG3:.+]]: index, %[[ARG4:.+]]: index, %[[ARG5:.+]]: index) -> tensor<?x?x?x?x?x?x?xf32> {
471471
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
472472
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
@@ -491,17 +491,17 @@ func.func @linalg_transpose_tensor_pack_multiple_tiles(%arg0: tensor<?x32x128xbf
491491
%dim = tensor.dim %arg0, %c0 : tensor<?x32x128xbf16>
492492

493493
%0 = tensor.empty(%dim) : tensor<32x128x?xbf16>
494-
%transposed = linalg.transpose
495-
ins(%arg0 : tensor<?x32x128xbf16>)
496-
outs(%0 : tensor<32x128x?xbf16>)
494+
%transposed = linalg.transpose
495+
ins(%arg0 : tensor<?x32x128xbf16>)
496+
outs(%0 : tensor<32x128x?xbf16>)
497497
permutation = [1, 2, 0]
498498

499499
%2 = tensor.empty(%dim) : tensor<32x?x64x16x2xbf16>
500-
%pack = tensor.pack %transposed
501-
padding_value(%cst : bf16)
502-
outer_dims_perm = [0, 2, 1]
503-
inner_dims_pos = [2, 1]
504-
inner_tiles = [16, 2]
500+
%pack = tensor.pack %transposed
501+
padding_value(%cst : bf16)
502+
outer_dims_perm = [0, 2, 1]
503+
inner_dims_pos = [2, 1]
504+
inner_tiles = [16, 2]
505505
into %2 : tensor<32x128x?xbf16> -> tensor<32x?x64x16x2xbf16>
506506
return %pack : tensor<32x?x64x16x2xbf16>
507507
}
@@ -513,11 +513,11 @@ func.func @linalg_transpose_tensor_pack_multiple_tiles(%arg0: tensor<?x32x128xbf
513513
// CHECK: %[[DIM:.+]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x32x128xbf16>
514514
// CHECK: %[[VAL0:.+]] = affine.apply #[[$MAP]]()[%[[DIM]]]
515515
// CHECK: %[[VAL1:.+]] = tensor.empty(%[[VAL0]]) : tensor<32x?x64x16x2xbf16>
516-
// CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]]
517-
// CHECK-SAME: padding_value(%[[CST]] : bf16)
518-
// CHECK-SAME: outer_dims_perm = [1, 0, 2]
519-
// CHECK-SAME: inner_dims_pos = [0, 2]
520-
// CHECK-SAME: inner_tiles = [16, 2]
516+
// CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]]
517+
// CHECK-SAME: padding_value(%[[CST]] : bf16)
518+
// CHECK-SAME: outer_dims_perm = [1, 0, 2]
519+
// CHECK-SAME: inner_dims_pos = [0, 2]
520+
// CHECK-SAME: inner_tiles = [16, 2]
521521
// CHECK-SAME: into %[[VAL1]] : tensor<?x32x128xbf16> -> tensor<32x?x64x16x2xbf16>
522522
// CHECK: return %[[PACK]] : tensor<32x?x64x16x2xbf16>
523523
// CHECK: }
@@ -526,24 +526,24 @@ func.func @linalg_transpose_tensor_pack_multiple_tiles(%arg0: tensor<?x32x128xbf
526526

527527
func.func @linalg_transpose_tensor_unpack_fold(%arg0: tensor<1x1x4x16xi32>) -> tensor<16x4xi32> {
528528
%0 = tensor.empty() : tensor<1x1x16x4xi32>
529-
%transposed = linalg.transpose ins(%arg0 : tensor<1x1x4x16xi32>)
530-
outs(%0 : tensor<1x1x16x4xi32>)
529+
%transposed = linalg.transpose ins(%arg0 : tensor<1x1x4x16xi32>)
530+
outs(%0 : tensor<1x1x16x4xi32>)
531531
permutation = [1, 0, 3, 2]
532532
%1 = tensor.empty() : tensor<16x4xi32>
533-
%unpack = tensor.unpack %transposed
534-
outer_dims_perm = [0, 1]
535-
inner_dims_pos = [0, 1]
536-
inner_tiles = [16, 4] into
533+
%unpack = tensor.unpack %transposed
534+
outer_dims_perm = [0, 1]
535+
inner_dims_pos = [0, 1]
536+
inner_tiles = [16, 4] into
537537
%1 : tensor<1x1x16x4xi32> -> tensor<16x4xi32>
538538
return %unpack : tensor<16x4xi32>
539539
}
540540
//CHECK-LABEL: func.func @linalg_transpose_tensor_unpack_fold(
541541
// CHECK-SAME: %[[ARG0:.+]]: tensor<1x1x4x16xi32>) -> tensor<16x4xi32> {
542542
// CHECK: %[[OUT:.+]] = tensor.empty() : tensor<16x4xi32>
543-
// CHECK: %[[UNPACK:.+]] = tensor.unpack %[[ARG0]]
544-
// CHECK-SAME: outer_dims_perm = [1, 0]
545-
// CHECK-SAME: inner_dims_pos = [1, 0]
546-
// CHECK-SAME: inner_tiles = [4, 16]
543+
// CHECK: %[[UNPACK:.+]] = tensor.unpack %[[ARG0]]
544+
// CHECK-SAME: outer_dims_perm = [1, 0]
545+
// CHECK-SAME: inner_dims_pos = [1, 0]
546+
// CHECK-SAME: inner_tiles = [4, 16]
547547
// CHEKC-SAME: into %[[OUT]] : tensor<1x1x4x16xi32> -> tensor<16x4xi32>
548548
// CHECK: return %[[UNPACK]] : tensor<16x4xi32>
549549
// CHECK: }
@@ -555,7 +555,7 @@ func.func @linalg_transpose_tensor_unpack_fold_dynamic_outer_dims_tile_dims_tile
555555
ins(%arg0 : tensor<?x?x?x?xf32>)
556556
outs(%transpose_dest : tensor<?x?x?x?xf32>)
557557
permutation = [1, 0, 3, 2]
558-
558+
559559
%unpack = tensor.unpack %transposed
560560
outer_dims_perm = [1, 0]
561561
inner_dims_pos = [0, 1]
@@ -565,7 +565,7 @@ func.func @linalg_transpose_tensor_unpack_fold_dynamic_outer_dims_tile_dims_tile
565565
}
566566
// CHECK: #[[$MAP:.+]] = affine_map<()[s0, s1] -> (s0 * s1)>
567567
// CHECK-LABEL: func.func @linalg_transpose_tensor_unpack_fold_dynamic_outer_dims_tile_dims_tile_sizes(
568-
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?xf32>, %[[ARG1:.+]]: tensor<?x?x?x?xf32>, %[[ARG2:.+]]: tensor<?x?xf32>,
568+
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?xf32>, %[[ARG1:.+]]: tensor<?x?x?x?xf32>, %[[ARG2:.+]]: tensor<?x?xf32>,
569569
// CHECK-SAME: %[[IDX1:.+]]: index, %[[IDX2:.+]]: index) -> tensor<?x?xf32> {
570570
// CHECK-DAG: %[[CST1:.+]] = arith.constant 1 : index
571571
// CHECK-DAG: %[[CST0:.+]] = arith.constant 0 : index
@@ -574,10 +574,10 @@ func.func @linalg_transpose_tensor_unpack_fold_dynamic_outer_dims_tile_dims_tile
574574
// CHECK-DAG: %[[AMAP0:.+]] = affine.apply #[[$MAP]]()[%[[DIM1]], %[[IDX2]]]
575575
// CHECK-DAG: %[[AMAP1:.+]] = affine.apply #[[$MAP]]()[%[[DIM0]], %[[IDX1]]]
576576
// CHECK: %[[OUT:.+]] = tensor.empty(%[[AMAP1]], %[[AMAP0]]) : tensor<?x?xf32>
577-
// CHECK: %[[UNPACK:.+]] = tensor.unpack %[[ARG0]]
578-
// CHECK-SAME: outer_dims_perm = [0, 1]
579-
// CHECK-SAME: inner_dims_pos = [1, 0]
580-
// CHECK-SAME: inner_tiles = [%[[IDX2]], %[[IDX1]]]
577+
// CHECK: %[[UNPACK:.+]] = tensor.unpack %[[ARG0]]
578+
// CHECK-SAME: outer_dims_perm = [0, 1]
579+
// CHECK-SAME: inner_dims_pos = [1, 0]
580+
// CHECK-SAME: inner_tiles = [%[[IDX2]], %[[IDX1]]]
581581
// CHECK-SAME: into %[[OUT]] : tensor<?x?x?x?xf32> -> tensor<?x?xf32>
582582
// CHECK: return %[[UNPACK]] : tensor<?x?xf32>
583583
// CHECK: }
@@ -602,10 +602,37 @@ func.func @tensor_unpack_linalg_transpose_fold(%arg0: tensor<56x57x1x64xf32>) ->
602602
// CHECK-LABEL: func.func @tensor_unpack_linalg_transpose_fold(
603603
// CHECK-SAME: %[[ARG0:.+]]: tensor<56x57x1x64xf32>) -> tensor<3648x56xf32> {
604604
// CHECK: %[[OUT:.+]] = tensor.empty() : tensor<3648x56xf32>
605-
// CHECK: %[[UNPACK:.+]] = tensor.unpack %[[ARG0]]
606-
// CHECK-SAME: outer_dims_perm = [1, 0]
607-
// CHECK-SAME: inner_dims_pos = [1, 0]
608-
// CHECK-SAME: inner_tiles = [1, 64]
605+
// CHECK: %[[UNPACK:.+]] = tensor.unpack %[[ARG0]]
606+
// CHECK-SAME: outer_dims_perm = [1, 0]
607+
// CHECK-SAME: inner_dims_pos = [1, 0]
608+
// CHECK-SAME: inner_tiles = [1, 64]
609609
// CHECK-SAME: into %[[OUT:.+]] : tensor<56x57x1x64xf32> -> tensor<3648x56xf32>
610610
// CHECK: return %[[UNPACK]] : tensor<3648x56xf32>
611611
// CHECK: }
612+
613+
// -----
614+
615+
func.func @tensor_padded_unpack_linalg_transpose_fold(%arg0: tensor<71x7x4x16x16xf32>) -> tensor<100x71x64xf32> {
616+
%0 = tensor.empty() : tensor<71x100x64xf32>
617+
%pack = tensor.unpack %arg0
618+
inner_dims_pos = [1, 2]
619+
inner_tiles = [16, 16]
620+
into %0 : tensor<71x7x4x16x16xf32> -> tensor<71x100x64xf32>
621+
622+
%1 = tensor.empty() : tensor<100x71x64xf32>
623+
%transposed = linalg.transpose
624+
ins(%pack : tensor<71x100x64xf32>)
625+
outs(%1 : tensor<100x71x64xf32>)
626+
permutation = [1, 0, 2]
627+
return %transposed : tensor<100x71x64xf32>
628+
}
629+
// CHECK-LABEL: func.func @tensor_padded_unpack_linalg_transpose_fold(
630+
// CHECK-SAME: %[[ARG0:.+]]: tensor<71x7x4x16x16xf32>) -> tensor<100x71x64xf32> {
631+
// CHECK: %[[OUT:.+]] = tensor.empty() : tensor<100x71x64xf32>
632+
// CHECK: %[[UNPACK:.+]] = tensor.unpack %[[ARG0]]
633+
// CHECK-SAME: outer_dims_perm = [1, 0, 2]
634+
// CHECK-SAME: inner_dims_pos = [0, 2]
635+
// CHECK-SAME: inner_tiles = [16, 16]
636+
// CHECK-SAME: into %[[OUT:.+]] : tensor<71x7x4x16x16xf32> -> tensor<100x71x64xf32>
637+
// CHECK: return %[[UNPACK]] : tensor<100x71x64xf32>
638+
// CHECK: }

0 commit comments

Comments
 (0)