Add MSCCL Support (#658)

* Add MSCCL support

* Add alignment and message size checking

* Fix nRanks checking, in-place and out-of-place tests and group call handling

* Fix hipGraph unit test

* Change MSCCL init warning to INFO

* Revise license info

[ROCm/rccl commit: adafc0f759]
This commit is contained in:
Ziyue Yang
2022-12-13 07:51:04 +08:00
committed by GitHub
parent bb27f5c8ba
commit 431c7d2cb9
38 changed files with 40923 additions and 11 deletions
+15
View File
@@ -1,6 +1,7 @@
/*************************************************************************
* Copyright (c) 2015-2022, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
* Modifications Copyright (c) Microsoft Corporation. Licensed under the MIT License.
*
* See LICENSE.txt for license information
************************************************************************/
@@ -12,6 +13,8 @@
#include "channel.h"
#include <assert.h>
#include "msccl/msccl_lifecycle.h"
__thread int ncclGroupDepth = 0; // depth of ncclGroupStart nesting
__thread ncclResult_t ncclGroupError = ncclSuccess;
__thread struct ncclComm* ncclGroupCommHead = nullptr;
@@ -100,6 +103,14 @@ exit:
return ret;
}
ncclResult_t ncclGroupStartInternal() {
ncclGroupDepth++;
if (mscclAvailable() && !mscclIsCaller()) {
NCCLCHECK(mscclGroupStart());
}
return ncclSuccess;
}
NCCL_API(ncclResult_t, ncclGroupEnd);
ncclResult_t ncclGroupEnd() {
ncclResult_t ret = ncclSuccess;
@@ -386,6 +397,10 @@ ncclResult_t ncclGroupEndInternal() {
goto exit;
}
if (mscclAvailable() && !mscclIsCaller()) {
NCCLCHECK(mscclGroupEnd());
}
if ((--ncclGroupDepth) > 0) goto exit;
if ((ret = ncclGroupError) != ncclSuccess) goto fail;