[Mlir-commits] [mlir] [mlir][linalg] Add PackOp canonicalization pattern (PR #215785)
Hendrik Klug
llvmlistbot at llvm.org
Mon Aug 24 12:51:27 PDT 2026
================
@@ -6008,6 +6008,13 @@ LogicalResult PackOp::canonicalize(PackOp packOp, PatternRewriter &rewriter) {
}
}
+ // Fold pack(empty) to the destination tensor if no padding value is provided.
+ if (packOp.getSource().getDefiningOp<tensor::EmptyOp>() &&
----------------
HendrikKlug-synthara wrote:
for the few examples I could come up with, bufferization seems to work fine:
<details>
<summary>Users of packed empty</summary>
```mlir
#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
module {
func.func @init_from_empty() -> (tensor<4x8x8x32xf32>, tensor<4x8x8x32xf32>, tensor<4x8x8x32xf32>, tensor<64x128xf32>) {
%cst = arith.constant 1.000000e+00 : f32
%0 = tensor.empty() : tensor<64x128xf32>
%1 = tensor.empty() : tensor<4x8x8x32xf32>
// packing of empty tensor into empty tensor: will get folded into a single empty tensor with the packed shape (%1)
// 0% stays alive because its returned
%pack = linalg.pack %0 outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [8, 32] into %1 : tensor<64x128xf32> -> tensor<4x8x8x32xf32>
// bunch of users of the packed tensor
%2 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} outs(%1 : tensor<4x8x8x32xf32>) {
^bb0(%out: f32):
%6 = arith.addf %out, %cst : f32
linalg.yield %6 : f32
} -> tensor<4x8x8x32xf32>
%3 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} outs(%2 : tensor<4x8x8x32xf32>) {
^bb0(%out: f32):
%6 = arith.addf %out, %cst : f32
linalg.yield %6 : f32
} -> tensor<4x8x8x32xf32>
%4 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} outs(%1 : tensor<4x8x8x32xf32>) {
^bb0(%out: f32):
%6 = arith.addf %out, %cst : f32
linalg.yield %6 : f32
} -> tensor<4x8x8x32xf32>
%5 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} outs(%4 : tensor<4x8x8x32xf32>) {
^bb0(%out: f32):
%6 = arith.addf %out, %cst : f32
linalg.yield %6 : f32
} -> tensor<4x8x8x32xf32>
return %3, %5, %pack, %0 : tensor<4x8x8x32xf32>, tensor<4x8x8x32xf32>, tensor<4x8x8x32xf32>, tensor<64x128xf32>
}
}
// -----// IR Dump After CanonicalizerPass: canonicalize{cse-between-iterations=false max-iterations=10 max-num-rewrites=-1 region-simplify=normal test-convergence=false top-down=true} //----- //
#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
module {
func.func @init_from_empty() -> (tensor<4x8x8x32xf32>, tensor<4x8x8x32xf32>, tensor<4x8x8x32xf32>, tensor<64x128xf32>) {
%cst = arith.constant 1.000000e+00 : f32
%0 = tensor.empty() : tensor<64x128xf32>
%1 = tensor.empty() : tensor<4x8x8x32xf32>
%2 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} outs(%1 : tensor<4x8x8x32xf32>) {
^bb0(%out: f32):
%6 = arith.addf %out, %cst : f32
linalg.yield %6 : f32
} -> tensor<4x8x8x32xf32>
%3 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} outs(%2 : tensor<4x8x8x32xf32>) {
^bb0(%out: f32):
%6 = arith.addf %out, %cst : f32
linalg.yield %6 : f32
} -> tensor<4x8x8x32xf32>
%4 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} outs(%1 : tensor<4x8x8x32xf32>) {
^bb0(%out: f32):
%6 = arith.addf %out, %cst : f32
linalg.yield %6 : f32
} -> tensor<4x8x8x32xf32>
%5 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} outs(%4 : tensor<4x8x8x32xf32>) {
^bb0(%out: f32):
%6 = arith.addf %out, %cst : f32
linalg.yield %6 : f32
} -> tensor<4x8x8x32xf32>
return %3, %5, %1, %0 : tensor<4x8x8x32xf32>, tensor<4x8x8x32xf32>, tensor<4x8x8x32xf32>, tensor<64x128xf32>
}
}
// -----// IR Dump After OneShotBufferizePass: one-shot-bufferize //----- //
#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
module {
func.func @init_from_empty() -> (tensor<4x8x8x32xf32>, tensor<4x8x8x32xf32>, tensor<4x8x8x32xf32>, tensor<64x128xf32>) {
%cst = arith.constant 1.000000e+00 : f32
%alloc = memref.alloc() {alignment = 64 : i64} : memref<64x128xf32>
%0 = bufferization.to_tensor %alloc : memref<64x128xf32> to tensor<64x128xf32>
%alloc_0 = memref.alloc() {alignment = 64 : i64} : memref<4x8x8x32xf32>
%1 = bufferization.to_tensor %alloc_0 : memref<4x8x8x32xf32> to tensor<4x8x8x32xf32>
%alloc_1 = memref.alloc() {alignment = 64 : i64} : memref<4x8x8x32xf32>
linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} outs(%alloc_1 : memref<4x8x8x32xf32>) {
^bb0(%out: f32):
%4 = arith.addf %out, %cst : f32
linalg.yield %4 : f32
}
linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} outs(%alloc_1 : memref<4x8x8x32xf32>) {
^bb0(%out: f32):
%4 = arith.addf %out, %cst : f32
linalg.yield %4 : f32
}
%2 = bufferization.to_tensor %alloc_1 : memref<4x8x8x32xf32> to tensor<4x8x8x32xf32>
%alloc_2 = memref.alloc() {alignment = 64 : i64} : memref<4x8x8x32xf32>
linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} outs(%alloc_2 : memref<4x8x8x32xf32>) {
^bb0(%out: f32):
%4 = arith.addf %out, %cst : f32
linalg.yield %4 : f32
}
linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} outs(%alloc_2 : memref<4x8x8x32xf32>) {
^bb0(%out: f32):
%4 = arith.addf %out, %cst : f32
linalg.yield %4 : f32
}
%3 = bufferization.to_tensor %alloc_2 : memref<4x8x8x32xf32> to tensor<4x8x8x32xf32>
return %2, %3, %1, %0 : tensor<4x8x8x32xf32>, tensor<4x8x8x32xf32>, tensor<4x8x8x32xf32>, tensor<64x128xf32>
}
}
```
</details>
A slightly more interesting case which actually fails to bufferize currently:
```mlir
module {
func.func @write_then_pack(%arg0: f32, %arg1: tensor<4x8x8x32xf32>) -> (tensor<4x8x8x32xf32>, tensor<64x128xf32>) {
%0 = tensor.empty() : tensor<64x128xf32>
%1 = linalg.fill ins(%arg0 : f32) outs(%0 : tensor<64x128xf32>) -> tensor<64x128xf32>
%pack = linalg.pack %0 outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [8, 32] into %arg1 : tensor<64x128xf32> ->
tensor<4x8x8x32xf32>
return %pack, %1 : tensor<4x8x8x32xf32>, tensor<64x128xf32>
}
}
````
<details>
<summary>without empty pack canonicalization</summary>
```mlir
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
module {
func.func @write_then_pack(%arg0: f32, %arg1: tensor<4x8x8x32xf32>) -> (tensor<4x8x8x32xf32>, tensor<64x128xf32>) {
%0 = tensor.empty() : tensor<64x128xf32>
%1 = linalg.fill ins(%arg0 : f32) outs(%0 : tensor<64x128xf32>) -> tensor<64x128xf32>
%pack = linalg.pack %0 outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [8, 32] into %arg1 : tensor<64x128xf32> -> tensor<4x8x8x32xf32>
return %pack, %1 : tensor<4x8x8x32xf32>, tensor<64x128xf32>
}
}
/tmp/hjk/tmp.mlir:8:8: error: op was not bufferized
%p = linalg.pack %e outer_dims_perm = [1, 0] inner_dims_pos = [0, 1]
^
/tmp/hjk/tmp.mlir:8:8: note: see current operation: %3 = "linalg.pack"(%1, %arg1) <{inner_dims_pos = array<i64: 0, 1>, operandSegmentSizes = array<i32: 1, 1, 0, 0>, outer_dims_perm = array<i64: 1, 0>, static_inner_tiles = array<i64: 8, 32>}> : (tensor<64x128xf32>, tensor<4x8x8x32xf32>) -> tensor<4x8x8x32xf32>
// -----// IR Dump After OneShotBufferizePass Failed (one-shot-bufferize) //----- //
module {
func.func @write_then_pack(%arg0: f32, %arg1: tensor<4x8x8x32xf32>) -> (tensor<4x8x8x32xf32>, tensor<64x128xf32>) {
%alloc = memref.alloc() {alignment = 64 : i64} : memref<64x128xf32>
%0 = bufferization.to_tensor %alloc : memref<64x128xf32> to tensor<64x128xf32>
linalg.fill ins(%arg0 : f32) outs(%alloc : memref<64x128xf32>)
%1 = bufferization.to_tensor %alloc : memref<64x128xf32> to tensor<64x128xf32>
%pack = linalg.pack %0 outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [8, 32] into %arg1 : tensor<64x128xf32> -> tensor<4x8x8x32xf32>
return %pack, %1 : tensor<4x8x8x32xf32>, tensor<64x128xf32>
}
}
```
</details>
<details>
<summary>with empty pack canonicalization</summary>
```mlir
// -----// IR Dump After CanonicalizerPass: canonicalize{cse-between-iterations=false max-iterations=10 max-num-rewrites=-1 region-simplify=normal test-convergence=false top-down=true} //----- //
module {
func.func @write_then_pack(%arg0: f32, %arg1: tensor<4x8x8x32xf32>) -> (tensor<4x8x8x32xf32>, tensor<64x128xf32>) {
%0 = tensor.empty() : tensor<64x128xf32>
%1 = linalg.fill ins(%arg0 : f32) outs(%0 : tensor<64x128xf32>) -> tensor<64x128xf32>
return %arg1, %1 : tensor<4x8x8x32xf32>, tensor<64x128xf32>
}
}
// -----// IR Dump After OneShotBufferizePass: one-shot-bufferize{allow-return-allocs-from-loops=false allow-unknown-ops=false analysis-fuzzer-seed=0 analysis-heuristic=bottom-up buffer-alignment=64 bufferize-function-boundaries=false check-parallel-regions=true copy-before-write=false dump-alias-sets=false function-boundary-type-conversion=infer-layout-map must-infer-memory-space=false print-conflicts=false test-analysis-only=false unknown-type-conversion=fully-dynamic-layout-map use-encoding-for-memory-space=false} //----- //
module {
func.func @write_then_pack(%arg0: f32, %arg1: tensor<4x8x8x32xf32>) -> (tensor<4x8x8x32xf32>, tensor<64x128xf32>) {
%alloc = memref.alloc() {alignment = 64 : i64} : memref<64x128xf32>
linalg.fill ins(%arg0 : f32) outs(%alloc : memref<64x128xf32>)
%0 = bufferization.to_tensor %alloc : memref<64x128xf32> to tensor<64x128xf32>
return %arg1, %0 : tensor<4x8x8x32xf32>, tensor<64x128xf32>
}
}
module {
func.func @write_then_pack(%arg0: f32, %arg1: tensor<4x8x8x32xf32>) -> (tensor<4x8x8x32xf32>, tensor<64x128xf32>) {
%alloc = memref.alloc() {alignment = 64 : i64} : memref<64x128xf32>
linalg.fill ins(%arg0 : f32) outs(%alloc : memref<64x128xf32>)
%0 = bufferization.to_tensor %alloc : memref<64x128xf32> to tensor<64x128xf32>
return %arg1, %0 : tensor<4x8x8x32xf32>, tensor<64x128xf32>
}
}
```
</details>
https://github.com/llvm/llvm-project/pull/215785
More information about the Mlir-commits
mailing list