diff --git a/src/collectives/device/common.h b/src/collectives/device/common.h index de1e6d84ec..aa4d8f7e32 100644 --- a/src/collectives/device/common.h +++ b/src/collectives/device/common.h @@ -11,7 +11,7 @@ #include "collectives.h" #include "devcomm.h" -#if defined(__gfx908__) +#if defined(__gfx908__) || defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__) #define COLL_UNROLL 2 #else #define COLL_UNROLL 4 diff --git a/src/collectives/device/sendrecv.h b/src/collectives/device/sendrecv.h index 5c5b5b72ed..7e3cff2664 100644 --- a/src/collectives/device/sendrecv.h +++ b/src/collectives/device/sendrecv.h @@ -58,8 +58,14 @@ struct RunWork { } #endif +#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__) + reduceCopy + (tid, nthreads, 0, nullptr, false, 1, &buff, 1, &recvBuff, count); +#else reduceCopy (tid, nthreads, 0, nullptr, false, 1, &buff, 1, &recvBuff, count); +#endif + #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_PRIM_SIMPLE_REDUCE_OR_COPY_MULTI_EXIT) if (isNpKitThread) { NpKit::CollectGpuEvent(NPKIT_EVENT_PRIM_SIMPLE_REDUCE_OR_COPY_MULTI_EXIT, count*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(), @@ -204,7 +210,7 @@ struct RunWork { } else { #if defined(__gfx90a__) runRecv>(tid, nthreads, group, args); -#elif defined(__gfx908__) +#elif defined(__gfx908__) || defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__) runRecv>(tid, nthreads, group, args); #else runRecv>(tid, nthreads, group, args); @@ -216,7 +222,7 @@ struct RunWork { } else { #if defined(__gfx90a__) runSend>(tid, nthreads, group, args); -#elif defined(__gfx908__) +#elif defined(__gfx908__) || defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__) runSend>(tid, nthreads, group, args); #else runSend>(tid, nthreads, group, args);