diff --git a/projects/rccl-tests/src/common.cu b/projects/rccl-tests/src/common.cu index 23d884cb9b..401ba46c2b 100644 --- a/projects/rccl-tests/src/common.cu +++ b/projects/rccl-tests/src/common.cu @@ -32,7 +32,7 @@ const char *test_typenames[ncclNumTypes] = {"char", "int", "half", "float", "dou #endif ncclRedOp_t test_ops[ncclNumOps] = {ncclSum, ncclProd, ncclMax, ncclMin}; const char *test_opnames[ncclNumOps] = {"sum", "prod", "max", "min"}; -const char *test_memorytypes[nccl_NUM_MTYPES] = {"coarse", "fine", "host"}; +const char *test_memorytypes[nccl_NUM_MTYPES] = {"coarse", "fine", "host", "managed"}; thread_local int is_main_thread = 0; @@ -655,6 +655,11 @@ testResult_t AllocateBuffs(void **sendbuff, size_t sendBytes, void **recvbuff, s HIPCHECK(hipHostMalloc(recvbuff, nbytes)); HIPCHECK(hipHostMalloc(expected, recvBytes)); } + else if (memorytype == ncclManaged) { + HIPCHECK(hipMallocManaged(sendbuff, nbytes)); + HIPCHECK(hipMallocManaged(recvbuff, nbytes)); + HIPCHECK(hipMallocManaged(expected, recvBytes)); + } else { HIPCHECK(hipMalloc(sendbuff, nbytes)); HIPCHECK(hipMalloc(recvbuff, nbytes)); @@ -813,7 +818,7 @@ int main(int argc, char* argv[]) { "[-d,--datatype ] \n\t" "[-r,--root ] \n\t" "[-z,--blocking <0/1>] \n\t" - "[-y,--memory_type ] \n\t" + "[-y,--memory_type ] \n\t" "[-s,--stress_cycles ] \n\t" "[-u,--cumask ] \n\t" "[-h,--help]\n", diff --git a/projects/rccl-tests/src/common.h b/projects/rccl-tests/src/common.h index 8de2efaa4b..a498cce8e6 100644 --- a/projects/rccl-tests/src/common.h +++ b/projects/rccl-tests/src/common.h @@ -212,7 +212,8 @@ extern const char *test_opnames[ncclNumOps]; typedef enum { ncclCoarse = 0, ncclFine = 1, ncclHost = 2, - nccl_NUM_MTYPES = 3 } ncclMemoryType_t; + ncclManaged = 3, + nccl_NUM_MTYPES = 4 } ncclMemoryType_t; extern const char *test_memorytypes[nccl_NUM_MTYPES]; static int ncclstringtotype(char *str) {