[Mlir-commits] [mlir] [mlir][bufferization] Support matching permutation maps in Linalg analysis (PR #226079)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Sep 24 02:02:28 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Ziliang Zhang (ziliangzl)
<details>
<summary>Changes</summary>
Generalize `LinalgOpInterface::bufferizesToElementwiseAccess` to recognize participating tensor operands with identical permutation indexing maps, instead of requiring identity maps.
This enables in-place bufferization after loop interchange. Projected and mismatching maps, reduction iterators, and sparse operands remain conservatively rejected.
The tests cover the analysis decision as well as end-to-end bufferization without an allocation or copy.
---
Full diff: https://github.com/llvm/llvm-project/pull/226079.diff
3 Files Affected:
- (modified) mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp (+7-4)
- (modified) mlir/test/Dialect/Linalg/one-shot-bufferize-analysis.mlir (+123)
- (modified) mlir/test/Dialect/Linalg/one-shot-bufferize.mlir (+33)
``````````diff
diff --git a/mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp
index 28945642da099..f969d5b73ea25 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -123,10 +123,12 @@ struct LinalgOpInterface
if (linalgOp.getNumLoops() != linalgOp.getNumParallelLoops())
return false;
- // All index maps of tensors must be identity maps.
+ // All indexing maps of participating tensors must be the same
+ // permutation.
SmallVector<AffineMap> indexingMaps = linalgOp.getIndexingMapsArray();
assert(linalgOp->getNumOperands() == indexingMaps.size() &&
"unexpected number of indexing maps");
+ AffineMap commonIndexingMap;
for (auto [operand, map] :
llvm::zip(linalgOp->getOpOperands(), indexingMaps)) {
// Non-tensors do not participate in bufferization, so they can be
@@ -136,10 +138,11 @@ struct LinalgOpInterface
// Only consider operands in `opOperands`.
if (!llvm::is_contained(opOperands, &operand))
continue;
- // TODO: This could be generalized to other indexing maps. (All indexing
- // must be the same.)
- if (!map.isIdentity())
+ if (!map.isPermutation())
return false;
+ if (commonIndexingMap && commonIndexingMap != map)
+ return false;
+ commonIndexingMap = map;
}
return true;
diff --git a/mlir/test/Dialect/Linalg/one-shot-bufferize-analysis.mlir b/mlir/test/Dialect/Linalg/one-shot-bufferize-analysis.mlir
index a0922bdfcfbe4..e255b4cb1ce0f 100644
--- a/mlir/test/Dialect/Linalg/one-shot-bufferize-analysis.mlir
+++ b/mlir/test/Dialect/Linalg/one-shot-bufferize-analysis.mlir
@@ -97,3 +97,126 @@ func.func @elementwise_access_regression(%arg0: i32, %arg2: tensor<32x1xf32>, %a
}
return
}
+
+// -----
+
+#bias = affine_map<(d0, d1) -> (d0)>
+#transpose = affine_map<(d0, d1) -> (d1, d0)>
+
+// The two uses of %dst have the same permutation map. The projected map of
+// %bias does not participate in the conflict.
+// CHECK-LABEL: func @matching_permutation_maps_bufferize_inplace(
+// CHECK: linalg.generic
+// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]}
+func.func @matching_permutation_maps_bufferize_inplace(
+ %bias: tensor<16xf32>,
+ %dst: tensor<8x16xf32> {bufferization.writable = true})
+ -> tensor<8x16xf32> {
+ %0 = linalg.generic {
+ indexing_maps = [#bias, #transpose, #transpose],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%bias, %dst : tensor<16xf32>, tensor<8x16xf32>)
+ outs(%dst : tensor<8x16xf32>) {
+ ^bb0(%bias_elem: f32, %old: f32, %out: f32):
+ %1 = arith.addf %old, %bias_elem : f32
+ linalg.yield %1 : f32
+ } -> tensor<8x16xf32>
+ return %0 : tensor<8x16xf32>
+}
+
+// -----
+
+#identity = affine_map<(d0, d1) -> (d0, d1)>
+#transpose = affine_map<(d0, d1) -> (d1, d0)>
+
+// CHECK-LABEL: func @mismatching_permutation_maps(
+// CHECK: linalg.generic
+// CHECK-SAME: {__inplace_operands_attr__ = ["true", "false"]}
+func.func @mismatching_permutation_maps(
+ %arg0: tensor<4x4xf32> {bufferization.writable = true})
+ -> tensor<4x4xf32> {
+ %0 = linalg.generic {
+ indexing_maps = [#transpose, #identity],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%arg0 : tensor<4x4xf32>) outs(%arg0 : tensor<4x4xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %1 = arith.addf %in, %out : f32
+ linalg.yield %1 : f32
+ } -> tensor<4x4xf32>
+ return %0 : tensor<4x4xf32>
+}
+
+// -----
+
+#identity = affine_map<(d0, d1) -> (d0, d1)>
+#projected = affine_map<(d0, d1) -> (d1)>
+
+// Matching maps are insufficient when they are not permutations. Because
+// #projected drops d0, multiple iterations access the same element of %state.
+// CHECK-LABEL: func @reject_matching_projected_maps(
+// CHECK: linalg.generic
+// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "false"]}
+func.func @reject_matching_projected_maps(
+ %source: tensor<4x4xf32>,
+ %state: tensor<4xf32> {bufferization.writable = true})
+ -> tensor<4xf32> {
+ %0 = linalg.generic {
+ indexing_maps = [#identity, #projected, #projected],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%source, %state : tensor<4x4xf32>, tensor<4xf32>)
+ outs(%state : tensor<4xf32>) {
+ ^bb0(%source_elem: f32, %state_elem: f32, %out: f32):
+ %1 = arith.addf %source_elem, %state_elem : f32
+ linalg.yield %1 : f32
+ } -> tensor<4xf32>
+ return %0 : tensor<4xf32>
+}
+
+// -----
+
+#transpose = affine_map<(d0, d1) -> (d1, d0)>
+
+// This is intentionally not a semantic reduction. The reduction iterator
+// verifies that matching permutation maps are rejected unless all loops are
+// parallel.
+// CHECK-LABEL: func @reject_reduction_iterator_with_matching_permutation_maps(
+// CHECK: linalg.generic
+// CHECK-SAME: {__inplace_operands_attr__ = ["true", "false"]}
+func.func @reject_reduction_iterator_with_matching_permutation_maps(
+ %arg0: tensor<4x4xf32> {bufferization.writable = true})
+ -> tensor<4x4xf32> {
+ %0 = linalg.generic {
+ indexing_maps = [#transpose, #transpose],
+ iterator_types = ["parallel", "reduction"]}
+ ins(%arg0 : tensor<4x4xf32>) outs(%arg0 : tensor<4x4xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %1 = arith.addf %in, %out : f32
+ linalg.yield %1 : f32
+ } -> tensor<4x4xf32>
+ return %0 : tensor<4x4xf32>
+}
+
+// -----
+
+#sparse = #sparse_tensor.encoding<{
+ map = (d0, d1) -> (d0 : dense, d1 : compressed)
+}>
+#transpose = affine_map<(d0, d1) -> (d1, d0)>
+
+// CHECK-LABEL: func @matching_permutation_maps_with_sparse_operand(
+// CHECK: linalg.generic
+// CHECK-SAME: {__inplace_operands_attr__ = ["true", "false"]}
+func.func @matching_permutation_maps_with_sparse_operand(
+ %arg0: tensor<4x4xf32, #sparse> {bufferization.writable = true})
+ -> tensor<4x4xf32, #sparse> {
+ %0 = linalg.generic {
+ indexing_maps = [#transpose, #transpose],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%arg0 : tensor<4x4xf32, #sparse>)
+ outs(%arg0 : tensor<4x4xf32, #sparse>) {
+ ^bb0(%in: f32, %out: f32):
+ %1 = arith.addf %in, %out : f32
+ linalg.yield %1 : f32
+ } -> tensor<4x4xf32, #sparse>
+ return %0 : tensor<4x4xf32, #sparse>
+}
diff --git a/mlir/test/Dialect/Linalg/one-shot-bufferize.mlir b/mlir/test/Dialect/Linalg/one-shot-bufferize.mlir
index 2d6e10b123159..8fff7df233863 100644
--- a/mlir/test/Dialect/Linalg/one-shot-bufferize.mlir
+++ b/mlir/test/Dialect/Linalg/one-shot-bufferize.mlir
@@ -10,6 +10,39 @@
// TODO: Some test cases from this file should be moved to other dialects.
+#bias = affine_map<(d0, d1) -> (d0)>
+#transpose = affine_map<(d0, d1) -> (d1, d0)>
+
+// CHECK-LABEL: func @interchanged_bias_add_inplace(
+// CHECK-SAME: %[[BIAS:.*]]: memref<16xf32
+// CHECK-SAME: %[[DST:.*]]: memref<8x16xf32
+// CHECK-NOT: memref.alloc
+// CHECK-NOT: memref.copy
+// CHECK: linalg.generic
+// CHECK-SAME: ins(%[[BIAS]], %[[DST]]
+// CHECK-SAME: outs(%[[DST]]
+func.func @interchanged_bias_add_inplace(
+ %bias: tensor<16xf32>,
+ %dst: tensor<8x16xf32> {bufferization.writable = true})
+ -> tensor<8x16xf32> {
+ // This indexing-map configuration can be produced by interchanging the
+ // loops of an identity-mapped in-place elementwise update. The projected
+ // map of %bias does not participate in the conflict between the two uses of
+ // %dst.
+ %0 = linalg.generic {
+ indexing_maps = [#bias, #transpose, #transpose],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%bias, %dst : tensor<16xf32>, tensor<8x16xf32>)
+ outs(%dst : tensor<8x16xf32>) {
+ ^bb0(%bias_elem: f32, %old: f32, %out: f32):
+ %1 = arith.addf %old, %bias_elem : f32
+ linalg.yield %1 : f32
+ } -> tensor<8x16xf32>
+ return %0 : tensor<8x16xf32>
+}
+
+// -----
+
// CHECK-LABEL: func private @fill_inplace(
// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<?xf32, strided<[?], offset: ?>>
// CHECK-NO-LAYOUT-MAP-LABEL: func private @fill_inplace(%{{.*}}: memref<?xf32>) {
``````````
</details>
https://github.com/llvm/llvm-project/pull/226079
More information about the Mlir-commits
mailing list