llvm / llvm/llvm-project

[mlir][nvgpu] pipeline_shared_memory_copies mishandles non-shared vector.transfer_write destinations

Open
#213,221 1 comment 0 reactions 0 assignees View on GitHub
mlir
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.