Broadcasts to Shared Memory on GPU Runs in Serial
- Dominant language
- Python
- Stars
- 593
- Forks
- 163
- Avg merge
- 2d 23h
- Merged PRs (30d)
- 60
Description
**Describe the bug**
When running code on a GPU, if you have a block of shared memory and you broadcast a variable to it, the generated CUDA assigns in serial. In my reproducer, the performance and behavior is identical, but I've encountered a case where this bug led to slower code that was incorrect.
**To Reproduce**
Steps to reproduce the behavior:
```py
@dace.program
def reproduce():
value: dace.float32 = 0
for i, j in dace.map[0:I:I1, 0:J:J1]:
smem = dace.ndarray((I1, J1), dtype=dace.float32, storage=dace.StorageType.GPU_Shared)
smem[:] = 0.0
for i1, j1 in dace.map[0:I1, 0:J1]:
value += smem[i1, j1]
return value
sdfg = reproduce.to_sdfg()
sdfg.apply_gpu_transformations()
blockentry = find_map_by_param(sdfg, 'j1')
blockentry.map.schedule = dace.ScheduleType.GPU_ThreadBlock
print(sdfg())
```
This code outputs the following CUDA for `smem[:] = 0`:
```cuda
__shared__ float smem[1024];
{
for (auto __i0 = 0; __i0 < 32; __i0 += 1) {
for (auto __i1 = 0; __i1 < 32; __i1 += 1) {
{
float __out;
///////////////////
// Tasklet code (assign_13_8)
__out = 0.0;
///////////////////
smem[((32 * __i0) + __i1)] = __out;
}
}
}
}
```
**Expected behavior**
Since the outermost loop is a `GPU_Device` loop, the broadcast should be a `GPU_ThreadBlock` loop. Using the following Python code:
```py
@dace.program
def correct():
value: dace.float32 = 0
for i, j in dace.map[0:I:I1, 0:J:J1]:
smem = dace.ndarray((I1, J1), dtype=dace.float32, storage=dace.StorageType.GPU_Shared)
for k, l in dace.map[0:I1, 0:J1]:
smem[k, l] = 0.0
for i1, j1 in dace.map[0:I1, 0:J1]:
value += smem[i1, j1]
return value
sdfg = correct.to_sdfg()
sdfg.apply_gpu_transformations()
blockentry = find_map_by_param(sdfg, 'j1')
blockentry.map.schedule = dace.ScheduleType.GPU_ThreadBlock
blockentry = find_map_by_param(sdfg, 'l')
blockentry.map.schedule = dace.ScheduleType.GPU_ThreadBlock
print(sdfg())
```
The correct CUDA is generated:
```cuda
__shared__ float smem[1024];
{
{
{
int l = threadIdx.x;
int k = threadIdx.y;
{
{
{
float __out;
///////////////////
// Tasklet code (assign_28_12)
__out = 0.0;
///////////////////
smem[((32 * k) + l)] = __out;
}
}
}
}
}
}
```
**Screenshots**
If applicable, add screenshots to help explain your problem.
**Desktop (please complete the following information):**
- OS: RHEL 7.9
- DaCe Version 0.14.2
Contributor guide
Research direction
Start with the reproducer and the call to apply_gpu_transformations, then inspect how the GPU_Device and GPU_ThreadBlock schedules are assigned for the broadcast to GPU shared memory. Compare the generated CUDA with the expected parallel version. Done means the broadcast uses a thread-block loop rather than serial nested loops and preserves the reported behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers, hpc
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100