[Mlir-commits] [mlir] [mlir][nvgpu] Fix crash in shared memory store detection (PR #213222)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Vaisman

<details>
<summary>Changes</summary>

`isStoreToShared` incorrectly used `||` instead of `&&` when checking
whether a `vector.transfer_write` targets shared memory.

This caused tensor destinations to hit an assertion failure when
assertions are enabled, and memrefs in the default address space to be
treated as shared-memory stores.

Use `&&` so only memrefs in the shared memory address space are
recognized as shared-memory stores.

Testing:
- Add regression tests for tensor destinations and default-address-space
  memrefs. Both fail without the fix.
- check-mlir

Fixes #<!-- -->213221

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


2 Files Affected:

- (modified) mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp (+1-1) 
- (modified) mlir/test/Dialect/NVGPU/transform-pipeline-shared.mlir (+76) 


``````````diff
diff --git a/mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp b/mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp
index 4e6b8ea43e698..039099c0d387c 100644
--- a/mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp
+++ b/mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp
@@ -166,7 +166,7 @@ static bool isStoreToShared(Operation *op, Value v) {
     return false;
 
   auto storeType = dyn_cast<MemRefType>(store.getBase().getType());
-  return storeType || hasSharedMemorySpace(storeType);
+  return storeType && hasSharedMemorySpace(storeType);
 }
 
 /// Returns true if the operation is a load from the default memory space the
diff --git a/mlir/test/Dialect/NVGPU/transform-pipeline-shared.mlir b/mlir/test/Dialect/NVGPU/transform-pipeline-shared.mlir
index e959949babd9e..a027ee4b7e603 100644
--- a/mlir/test/Dialect/NVGPU/transform-pipeline-shared.mlir
+++ b/mlir/test/Dialect/NVGPU/transform-pipeline-shared.mlir
@@ -184,3 +184,79 @@ module attributes {transform.with_named_sequence} {
     transform.yield
   }
 }
+
+// -----
+
+// Verify that tensor-backed vector.transfer_write operations are not treated
+// as stores to shared memory.
+func.func @tensor_transfer_write(
+    %global: memref<16xf32>,
+    %dest: tensor<16xf32>) -> tensor<16xf32> {
+  %c0 = arith.constant 0 : index
+  %c4 = arith.constant 4 : index
+  %c16 = arith.constant 16 : index
+  %c0f = arith.constant 0.0 : f32
+
+  // expected-error @below {{no shared memory copy}}
+  %result = scf.for %i = %c0 to %c16 step %c4
+      iter_args(%tensor = %dest) -> tensor<16xf32> {
+    %value = vector.transfer_read %global[%i], %c0f
+      : memref<16xf32>, vector<4xf32>
+    %updated = vector.transfer_write %value, %tensor[%i]
+      : vector<4xf32>, tensor<16xf32>
+    scf.yield %updated : tensor<16xf32>
+  }
+
+  return %result : tensor<16xf32>
+}
+
+!t = !transform.any_op
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(
+      %arg0: !t {transform.readonly}) {
+    %loop = transform.structured.match ops{["scf.for"]} in %arg0
+      : (!t) -> !t
+    transform.nvgpu.pipeline_shared_memory_copies
+        failures(propagate) %loop {depth = 2, peel_epilogue}
+        : (!t) -> !t
+    transform.yield
+  }
+}
+
+// -----
+
+// Verify that vector.transfer_write operations targeting the default memory
+// space are not treated as stores to shared memory.
+func.func @default_memory_transfer_write(
+    %source: memref<16xf32>,
+    %dest: memref<16xf32>) {
+  %c0 = arith.constant 0 : index
+  %c4 = arith.constant 4 : index
+  %c16 = arith.constant 16 : index
+  %c0f = arith.constant 0.0 : f32
+
+  // expected-error @below {{no shared memory copy}}
+  scf.for %i = %c0 to %c16 step %c4 {
+    %value = vector.transfer_read %source[%i], %c0f
+      : memref<16xf32>, vector<4xf32>
+    vector.transfer_write %value, %dest[%i]
+      : vector<4xf32>, memref<16xf32>
+  }
+
+  return
+}
+
+!t = !transform.any_op
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(
+      %arg0: !t {transform.readonly}) {
+    %loop = transform.structured.match ops{["scf.for"]} in %arg0
+      : (!t) -> !t
+    transform.nvgpu.pipeline_shared_memory_copies
+        failures(propagate) %loop {depth = 2, peel_epilogue}
+        : (!t) -> !t
+    transform.yield
+  }
+}

``````````

</details>


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