spcl / spcl/dace

Scheduling and WCR Issues of 'Pure' reduce expansion on CUDA

Open
#292 1 comment 0 reactions 0 assignees View on GitHub
bug codegen
Dominant language
Python
Stars
593
Forks
163
Avg merge
2d 23h
Merged PRs (30d)
60

Description

Expanding a reduce node on a CUDA graph can lead to wrong behaviour, especially when dealing with reduce nodes nested inside another map.

1. If the maps' schedules of the expanded reduce nodes are not set to sequential manually, thread indices may appear inside the loop leading to out of bounds access
2. If the maps' attribute wcr_conflict of the expanded reduce nodes is not set to False manually, other CUDA errors may occur (719 in the example below)

```
import dace
import numpy as np

N = dace.symbol('N')
M = dace.symbol('M')

N.set(300); M.set(300)

@dace.program
def TEST(A: dace.float32[M,N]):
return dace.reduce(lambda a, b: max(a,b), A, axis=1, identity = 0)

@dace.program
def TEST2(A: dace.float32[M,N]):
tmp_out = np.ndarray([M], dace.float32)
for i in dace.map[0:M]:
tmp_out[i] = dace.reduce(lambda a,b: max(a,b), A[i,:], identity=0)
return tmp_out

A = np.random.rand(M.get(), N.get()).astype(np.float32)

if __name__ == '__main__':
sdfg1 = TEST.to_sdfg()
sdfg2 = TEST2.to_sdfg()
sdfg1.apply_gpu_transformations()
sdfg2.apply_gpu_transformations()

'''
## sdfg2 does not work here
#return1 = sdfg1.compile()(A=A, N=N, M=M)
#return2 = sdfg2.compile()(A=A, N=N, M=M)

#print(np.linalg.norm(return1))
#print(np.linalg.norm(return2))
'''

sdfg1.expand_library_nodes()
sdfg2.expand_library_nodes()

# fix for getting sdfg2 to work
for sdfg in [sdfg2]:
for node in sdfg.nodes()[0].nodes():
if isinstance(node, dace.sdfg.nodes.NestedSDFG):
for state in node.sdfg.nodes():
for snode in state.nodes():
for e in state.out_edges(snode):
e.data.wcr_conflict = False
if isinstance(snode, dace.sdfg.nodes.MapEntry):
snode.schedule = dace.dtypes.ScheduleType.Sequential



return1 = sdfg1.compile()(A=A, N=N, M=M)
return2 = sdfg2.compile()(A=A, N=N, M=M)
print(np.linalg.norm(return1))
print(np.linalg.norm(return2))

```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.