Fix for CUDA codegen

dace/codegen/targets/cuda.py

@@ -1132,10 +1132,22 @@ def _emit_copy(self, state_id, src_node, src_storage, dst_node, dst_storage, dst
  func=funcname,
  type=dst_node.desc(sdfg).dtype.ctype,
  bdims=', '.join(_topy(self._block_dims)),
- is_async='true' if state_dfg.out_degree(dst_node) > 0 else 'true',
+ is_async='true' if state_dfg.out_degree(dst_node) > 0 else 'false',
  accum=accum,
  args=', '.join([src_expr] + _topy(src_strides) + [dst_expr] + custom_reduction +
  _topy(dst_strides) + _topy(copy_shape))), sdfg, state_id, [src_node, dst_node])
+ elif funcname == 'dace::SharedToGlobal1D':
+ # special case: use a new template struct that provides functions for copy and reduction
+ callsite_stream.write(
+ (' {func}<{type}, {bdims}, {copysize}, {is_async}>{accum}({args});').format(
+ func=funcname,
+ type=dst_node.desc(sdfg).dtype.ctype,
+ bdims=', '.join(_topy(self._block_dims)),
+ copysize=', '.join(_topy(copy_shape)),
+ is_async='true' if state_dfg.out_degree(dst_node) > 0 else 'false',
+ accum=accum or '::Copy',
+ args=', '.join([src_expr] + _topy(src_strides) + [dst_expr] + _topy(dst_strides) + custom_reduction)), sdfg,
+ state_id, [src_node, dst_node])
  else:
  callsite_stream.write(
  (' {func}<{type}, {bdims}, {copysize}, ' +
@@ -1145,7 +1157,7 @@ def _emit_copy(self, state_id, src_node, src_storage, dst_node, dst_storage, dst
  bdims=', '.join(_topy(self._block_dims)),
  copysize=', '.join(_topy(copy_shape)),
  dststrides=', '.join(_topy(dst_strides)),
- is_async='true' if state_dfg.out_degree(dst_node) > 0 else 'true',
+ is_async='true' if state_dfg.out_degree(dst_node) > 0 else 'false',
  accum=accum,
  args=', '.join([src_expr] + _topy(src_strides) + [dst_expr] + custom_reduction)), sdfg,
  state_id, [src_node, dst_node])

dace/runtime/include/dace/cuda/copy.cuh

@@ -736,58 +736,75 @@ namespace dace
  int COPY_XLEN, bool ASYNC>
  struct SharedToGlobal1D
  {
+ static constexpr int BLOCK_SIZE = BLOCK_WIDTH * BLOCK_HEIGHT * BLOCK_DEPTH;
+ static constexpr int TOTAL = COPY_XLEN;
+ static constexpr int WRITES = TOTAL / BLOCK_SIZE;
+ static constexpr int REM_WRITES = TOTAL % BLOCK_SIZE;
+
+ static DACE_DFI void Copy(const T *smem, int src_xstride, T *ptr, int dst_xstride)
+ {
+ if (!ASYNC)
+ __syncthreads();
+
+ // Linear thread ID
+ int ltid = GetLinearTID<BLOCK_WIDTH, BLOCK_HEIGHT, BLOCK_DEPTH>();
+
+ #pragma unroll
+ for (int i = 0; i < WRITES; ++i) {
+ *(ptr + (ltid + i * BLOCK_SIZE) * dst_xstride) =
+ *(smem + (ltid + i * BLOCK_SIZE) * src_xstride);
+ }
+
+ if (REM_WRITES != 0 && ltid < REM_WRITES) {
+ *(ptr + (ltid + WRITES*BLOCK_SIZE)* dst_xstride) =
+ *(smem + (ltid + WRITES * BLOCK_SIZE) * src_xstride);
+ }
+ }
+
  template <typename WCR>
- static DACE_DFI void Accum(const T *smem, int src_xstride, T *ptr, int DST_XSTRIDE, WCR wcr)
+ static DACE_DFI void Accum(const T *smem, int src_xstride, T *ptr, int dst_xstride, WCR wcr)
  {
  if (!ASYNC)
  __syncthreads();
 
  // Linear thread ID
  int ltid = GetLinearTID<BLOCK_WIDTH, BLOCK_HEIGHT, BLOCK_DEPTH>();
- constexpr int BLOCK_SIZE = BLOCK_WIDTH * BLOCK_HEIGHT * BLOCK_DEPTH;
- constexpr int TOTAL = COPY_XLEN;
- constexpr int WRITES = TOTAL / BLOCK_SIZE;
- constexpr int REM_WRITES = TOTAL % BLOCK_SIZE;
 
  #pragma unroll
  for (int i = 0; i < WRITES; ++i) {
  wcr_custom<T>::template reduce(
- wcr, ptr + (ltid + i * BLOCK_SIZE) * DST_XSTRIDE,
+ wcr, ptr + (ltid + i * BLOCK_SIZE) * dst_xstride,
  *(smem + (ltid + i * BLOCK_SIZE) * src_xstride));
  }
 
  if (REM_WRITES != 0) {
  if (ltid < REM_WRITES)
  wcr_custom<T>::template reduce(
- ptr + (ltid + WRITES * BLOCK_SIZE)* DST_XSTRIDE,
+ ptr + (ltid + WRITES * BLOCK_SIZE)* dst_xstride,
  *(smem + (ltid + WRITES * BLOCK_SIZE) * src_xstride));
  }
  }
 
  template <ReductionType REDTYPE>
- static DACE_DFI void Accum(const T *smem, int src_xstride, T *ptr, int DST_XSTRIDE)
+ static DACE_DFI void Accum(const T *smem, int src_xstride, T *ptr, int dst_xstride)
  {
  if (!ASYNC)
  __syncthreads();
 
  // Linear thread ID
  int ltid = GetLinearTID<BLOCK_WIDTH, BLOCK_HEIGHT, BLOCK_DEPTH>();
- constexpr int BLOCK_SIZE = BLOCK_WIDTH * BLOCK_HEIGHT * BLOCK_DEPTH;
- constexpr int TOTAL = COPY_XLEN;
- constexpr int WRITES = TOTAL / BLOCK_SIZE;
- constexpr int REM_WRITES = TOTAL % BLOCK_SIZE;
 
  #pragma unroll
  for (int i = 0; i < WRITES; ++i) {
  wcr_fixed<REDTYPE, T>::template reduce_atomic(
- ptr + (ltid + i * BLOCK_SIZE) * DST_XSTRIDE,
+ ptr + (ltid + i * BLOCK_SIZE) * dst_xstride,
  *(smem + (ltid + i * BLOCK_SIZE) * src_xstride));
  }
 
  if (REM_WRITES != 0) {
  if (ltid < REM_WRITES)
  wcr_fixed<REDTYPE, T>::template reduce_atomic(
- ptr + (ltid + WRITES*BLOCK_SIZE)* DST_XSTRIDE,
+ ptr + (ltid + WRITES*BLOCK_SIZE)* dst_xstride,
  *(smem + (ltid + WRITES * BLOCK_SIZE) * src_xstride));
  }
  }

tests/codegen/cuda_memcopy_test.py

@@ -0,0 +1,84 @@
+""" Tests code generation for array copy on GPU target. """
+import dace
+from dace.transformation.auto import auto_optimize
+
+import pytest
+import re
+
+# this test requires cupy module
+cp = pytest.importorskip("cupy")
+
+# initialize random number generator
+rng = cp.random.default_rng(42)
+
+
+@pytest.mark.gpu
+def test_gpu_shared_to_global_1D():
+ M = 32
+ N = dace.symbol('N')
+
+ @dace.program
+ def transpose_shared_to_global(A: dace.float64[M, N], B: dace.float64[N, M]):
+ for i in dace.map[0:N]:
+ local_gather = dace.define_local([M], A.dtype, storage=dace.StorageType.GPU_Shared)
+ for j in dace.map[0:M]:
+ local_gather[j] = A[j, i]
+ B[i, :] = local_gather
+
+
+ sdfg = transpose_shared_to_global.to_sdfg()
+ auto_optimize.apply_gpu_storage(sdfg)
+
+ size_M = M
+ size_N = 128
+
+ A = rng.random((size_M, size_N,))
+ B = rng.random((size_N, size_M,))
+
+ ref = A.transpose()
+
+ sdfg(A, B, N=size_N)
+ cp.allclose(ref, B)
+
+ code = sdfg.generate_code()[1].clean_code # Get GPU code (second file)
+ m = re.search('dace::SharedToGlobal1D<.+>::Copy', code)
+ assert m is not None
+
+
+@pytest.mark.gpu
+def test_gpu_shared_to_global_1D_accumulate():
+ M = 32
+ N = dace.symbol('N')
+
+ @dace.program
+ def transpose_and_add_shared_to_global(A: dace.float64[M, N], B: dace.float64[N, M]):
+ for i in dace.map[0:N]:
+ local_gather = dace.define_local([M], A.dtype, storage=dace.StorageType.GPU_Shared)
+ for j in dace.map[0:M]:
+ local_gather[j] = A[j, i]
+ local_gather[:] >> B(M, lambda x, y: x + y)[i, :]
+
+
+ sdfg = transpose_and_add_shared_to_global.to_sdfg()
+ auto_optimize.apply_gpu_storage(sdfg)
+
+ size_M = M
+ size_N = 128
+
+ A = rng.random((size_M, size_N,))
+ B = rng.random((size_N, size_M,))
+
+ ref = A.transpose() + B
+
+ sdfg(A, B, N=size_N)
+ cp.allclose(ref, B)
+
+ code = sdfg.generate_code()[1].clean_code # Get GPU code (second file)
+ m = re.search('dace::SharedToGlobal1D<.+>::template Accum', code)
+ assert m is not None
+
+
+if __name__ == '__main__':
+ test_gpu_shared_to_global_1D()
+ test_gpu_shared_to_global_1D_accumulate()
+