[mlir][nvgpu] pipeline_shared_memory_copies mishandles non-shared vector.transfer_write destinations
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
### Summary
`transform.nvgpu.pipeline_shared_memory_copies` incorrectly handles `vector.transfer_write`
operations whose destination is not in shared memory.
Two cases are affected:
- tensor destinations trigger an assertion failure;
- memrefs in the default memory space are treated as shared-memory stores.
### Expected behavior
Only `vector.transfer_write` operations targeting memrefs in the shared memory
address space should be recognized as shared-memory stores.
Tensor destinations and default-address-space memrefs should not be classified
as shared-memory stores. The transform should reject such loops with
a `no shared memory copy` diagnostic instead of crashing.
### Steps to reproduce
Save the following as `repro.mlir`:
```mlir
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
%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>
}
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(
%arg0: !transform.any_op {transform.readonly}) {
%loop = transform.structured.match ops{["scf.for"]} in %arg0
: (!transform.any_op) -> !transform.any_op
transform.nvgpu.pipeline_shared_memory_copies
failures(propagate) %loop {depth = 2, peel_epilogue}
: (!transform.any_op) -> !transform.any_op
transform.yield
}
}
```
### Run
```mlir-opt repro.mlir --transform-interpreter -canonicalize```
### Actual behavior
With assertions enabled, mlir-opt terminates with an assertion failure while
checking whether the tensor-backed vector.transfer_write targets shared memory.
Contributor guide
Research direction
Save the supplied reproducer as repro.mlir and run it with mlir-opt --transform-interpreter -canonicalize. Start at the transform.nvgpu.pipeline_shared_memory_copies entry point and trace how vector.transfer_write destinations are classified. Done means tensor and default-address-space destinations are not treated as shared-memory stores and the loop produces a no shared memory copy diagnostic instead of an assertion failure.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100