Stanley Tsang
d3fa257682
Adding multiprocess unit tests for collectives. To run, NCCL_COMM_ID=$HOSTNAME:12345 build/release/test/UnitTestsMultiProcess
79 строки
3.3 KiB
C++
79 строки
3.3 KiB
C++
/*************************************************************************
|
|
* Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
|
|
*
|
|
* See LICENSE.txt for license information
|
|
************************************************************************/
|
|
|
|
#ifndef TEST_COMBINEDCALLS_MULTI_PROCESS_HPP
|
|
#define TEST_COMBINEDCALLS_MULTI_PROCESS_HPP
|
|
|
|
#include "CorrectnessTest.hpp"
|
|
|
|
#include "test_AllGatherMultiProcess.hpp"
|
|
#include "test_AllReduceMultiProcess.hpp"
|
|
#include "test_BroadcastMultiProcess.hpp"
|
|
#include "test_ReduceMultiProcess.hpp"
|
|
#include "test_ReduceScatterMultiProcess.hpp"
|
|
|
|
namespace CorrectnessTests
|
|
{
|
|
class CombinedCallsMultiProcessCorrectnessTest : public MultiProcessCorrectnessTest
|
|
{
|
|
public:
|
|
void TestCombinedCalls(int rank, std::vector<Dataset*>& datasets, std::vector<ncclFunc_t> const& funcs)
|
|
{
|
|
SetUpPerProcess(rank, funcs, comms[rank], streams[rank], datasets);
|
|
|
|
if (numDevices > numDevicesAvailable) return;
|
|
|
|
Barrier barrier(rank, numDevices, std::atoi(getenv("NCCL_COMM_ID")));
|
|
|
|
// Compute expected results for each dataset in combined
|
|
int const root = 0;
|
|
AllGatherMultiProcessCorrectnessTest::ComputeExpectedResults(*datasets[0], barrier, rank, numDevices);
|
|
AllReduceMultiProcessCorrectnessTest::ComputeExpectedResults(*datasets[1], barrier, op, rank);
|
|
BroadcastMultiProcessCorrectnessTest::ComputeExpectedResults(*datasets[2], root, rank);
|
|
ReduceMultiProcessCorrectnessTest::ComputeExpectedResults(*datasets[3], barrier, op, root, rank);
|
|
ReduceScatterMultiProcessCorrectnessTest::ComputeExpectedResults(*datasets[4], barrier, op, rank);
|
|
|
|
size_t const byteCount = datasets[0]->NumBytes() / numDevices;
|
|
size_t const elemCount = numElements / numDevices;
|
|
|
|
ncclAllGather((int8_t *)datasets[0]->inputs[rank] + (rank * byteCount),
|
|
datasets[0]->outputs[rank], elemCount,
|
|
dataType, comms[rank], streams[rank]);
|
|
|
|
ncclAllReduce(datasets[1]->inputs[rank], datasets[1]->outputs[rank],
|
|
numElements, dataType, op, comms[rank], streams[rank]);
|
|
|
|
ncclBroadcast(datasets[2]->inputs[rank],
|
|
datasets[2]->outputs[rank],
|
|
numElements, dataType,
|
|
root, comms[rank], streams[rank]);
|
|
|
|
ncclReduce(datasets[3]->inputs[rank],
|
|
datasets[3]->outputs[rank],
|
|
numElements, dataType, op,
|
|
root, comms[rank], streams[rank]);
|
|
|
|
ncclReduceScatter(datasets[4]->inputs[rank],
|
|
(int8_t *)datasets[4]->outputs[rank] + (rank * byteCount),
|
|
elemCount, dataType, op,
|
|
comms[rank], streams[rank]);
|
|
|
|
// Wait for reduction to complete
|
|
HIP_CALL(hipStreamSynchronize(streams[rank]));
|
|
|
|
// Check results for each collective in the combined
|
|
for (int i = 0; i < 5; i++)
|
|
{
|
|
ValidateResults(*datasets[i], rank);
|
|
barrier.Wait();
|
|
datasets[i]->Release(rank);
|
|
}
|
|
}
|
|
};
|
|
}
|
|
|
|
#endif
|