[Mlir-commits] [mlir] [mlir][Linalg] Fix init operand read detection for linalg.map (PR #221330)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Bhavesh M (beamandala)

<details>
<summary>Changes</summary>

Fix `MapOp::payloadUsesValueFromOperand` to also check reads on `linalg::MapOp` init operands instead of assuming they're write only. This caused one-shot bufferization to omit a required copy when the payload reads the initial value. Without the copy the payload instead reads from an uninitialized buffer resulting in undefined behavior.

Fixes #<!-- -->213241

---
Full diff: https://github.com/llvm/llvm-project/pull/221330.diff


2 Files Affected:

- (modified) mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td (-1) 
- (modified) mlir/test/Dialect/Linalg/bufferize.mlir (+39) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
index 51bfcbec5d1a4..589b35e45a815 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
@@ -290,7 +290,6 @@ def MapOp : LinalgStructuredBase_Op<"map", [
     MutableOperandRange getDpsInitsMutable() { return getInitMutable(); }
 
     bool payloadUsesValueFromOperand(OpOperand * opOperand) {
-      if (isDpsInit(opOperand)) return false;
       return !getMatchingBlockArgument(opOperand).use_empty();
     }
 
diff --git a/mlir/test/Dialect/Linalg/bufferize.mlir b/mlir/test/Dialect/Linalg/bufferize.mlir
index 7b1af139f0f3e..b18448840bd8e 100644
--- a/mlir/test/Dialect/Linalg/bufferize.mlir
+++ b/mlir/test/Dialect/Linalg/bufferize.mlir
@@ -162,6 +162,45 @@ func.func @generic_with_init_tensor(%arg0: tensor<2x3x4xvector<3x4xi4>>,
 
 // -----
 
+// Test that an init operand that is read by the map payload is copied.
+
+// CHECK-LABEL: func @map_with_read_from_init(
+// CHECK-SAME:      %[[INIT:.*]]: tensor<4xf32>, %[[BUFFER:.*]]: memref<f32>)
+// CHECK:         %[[INIT_BUFFER:.*]] = bufferization.to_buffer %[[INIT]] : tensor<4xf32> to memref<4xf32>
+// CHECK:         %[[RESULT_BUFFER:.*]] = memref.alloc() {{.*}} : memref<4xf32>
+// CHECK:         memref.copy %[[INIT_BUFFER]], %[[RESULT_BUFFER]] : memref<4xf32> to memref<4xf32>
+// CHECK:         linalg.map outs(%[[RESULT_BUFFER]] : memref<4xf32>)
+// CHECK-NEXT:      (%[[OUT:.*]]: f32) {
+// CHECK:           memref.store %[[OUT]], %[[BUFFER]][] : memref<f32>
+func.func @map_with_read_from_init(%init: tensor<4xf32>,
+                                   %buffer: memref<f32>) -> tensor<4xf32> {
+  %cst = arith.constant 0.0 : f32
+  %mapped = linalg.map outs(%init : tensor<4xf32>) (%out: f32) {
+    memref.store %out, %buffer[] : memref<f32>
+    linalg.yield %cst : f32
+  }
+  return %mapped : tensor<4xf32>
+}
+
+// -----
+
+// Test that an init operand that isn't read by the map payload isn't copied.
+
+// CHECK-LABEL: func @map_without_read_from_init(
+// CHECK-SAME:      %{{.*}}: tensor<4xf32>)
+// CHECK:         %[[RESULT_BUFFER:.*]] = memref.alloc() {{.*}} : memref<4xf32>
+// CHECK-NOT:     memref.copy
+// CHECK:         linalg.map outs(%[[RESULT_BUFFER]] : memref<4xf32>)
+func.func @map_without_read_from_init(%init: tensor<4xf32>) -> tensor<4xf32> {
+  %cst = arith.constant 0.0 : f32
+  %mapped = linalg.map outs(%init : tensor<4xf32>) (%out: f32) {
+    linalg.yield %cst : f32
+  }
+  return %mapped : tensor<4xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @bufferize_fill(
 // CHECK-SAME:    %[[IN:.*]]: tensor<?xf32>
 func.func @bufferize_fill(%arg0: tensor<?xf32>) -> tensor<?xf32> {

``````````

</details>


https://github.com/llvm/llvm-project/pull/221330