/** * \brief NCCL implementation of \ref gpucomm_reduce. */ static int reduce(gpudata *src, size_t offsrc, gpudata *dest, size_t offdest, size_t count, int typecode, int opcode, int root, gpucomm *comm) { // need dummy init so that compiler shuts up ncclRedOp_t op = ncclNumOps; ncclDataType_t datatype = ncclNumTypes; gpudata *dst = NULL; int rank = 0; cuda_context *ctx; ASSERT_BUF(src); ASSERT_COMM(comm); GA_CHECK(get_rank(comm, &rank)); if (rank == root) { dst = dest; ASSERT_BUF(dest); } GA_CHECK(check_restrictions(src, offsrc, dst, offdest, count, typecode, opcode, comm, &datatype, &op)); ctx = comm->ctx; cuda_enter(ctx); // sync: wait till a write has finished (out of concurrent kernels) GA_CUDA_EXIT_ON_ERROR(ctx, cuda_wait(src, CUDA_WAIT_READ)); // sync: wait till a read/write has finished (out of concurrent kernels) if (rank == root) GA_CUDA_EXIT_ON_ERROR(ctx, cuda_wait(dest, CUDA_WAIT_WRITE)); // change stream of nccl ops to enable concurrency if (rank == root) NCCL_EXIT_ON_ERROR(ctx, ncclReduce((void *)(src->ptr + offsrc), (void *)(dest->ptr + offdest), count, datatype, op, root, comm->c, ctx->s)); else NCCL_EXIT_ON_ERROR(ctx, ncclReduce((void *)(src->ptr + offsrc), NULL, count, datatype, op, root, comm->c, ctx->s)); GA_CUDA_EXIT_ON_ERROR(ctx, cuda_record(src, CUDA_WAIT_READ)); if (rank == root) GA_CUDA_EXIT_ON_ERROR(ctx, cuda_record(dest, CUDA_WAIT_WRITE)); cuda_exit(ctx); return GA_NO_ERROR; }
PyObject * THCPModule_nccl_reduce(PyObject *self, PyObject *args) { HANDLE_TH_ERRORS PyObject *_inputs, *_outputs, *_streams; int root, op; if (!PyArg_ParseTuple(args, "OOOii", &_inputs, &_outputs, &_streams, &root, &op)) { THPUtils_invalidArguments(args, NULL, "nccl_reduce", 1, "(sequence[Tensor] inputs, sequence[Tensor]" " outputs, sequence[torch.cuda.Stream or None], int root, int op"); return NULL; } std::vector<at::Tensor> inputs = THPUtils_PySequence_to_TensorList(_inputs); std::vector<at::Tensor> outputs = THPUtils_PySequence_to_TensorList(_outputs); std::vector<THCStream*> streams = THPUtils_PySequence_to_THCStreamList(_streams); THPUtils_assert(inputs.size() == streams.size(), "number of streams is not equal to number of inputs"); // we can safely release GIL after this line, no python API used AutoNoGIL no_gil; _check_inputs(inputs, outputs, 1, 1); size_t len = inputs.size(); ncclDataType_t data_type = _get_data_type(inputs[0].type().ID()); int64_t count = inputs[0].numel(); std::lock_guard<std::mutex> lock(*(THCCachingAllocator_getCudaFreeMutex())); ncclComm_t *comm = _get_communicator(inputs); AutoGPU gpu_guard; #if defined(NCCL_MAJOR) && (NCCL_MAJOR >= 2) CHECK(ncclGroupStart()); #endif for (size_t i = 0; i < len; i++) { int device = inputs[i].get_device(); gpu_guard.setDevice(device); auto stream = (streams[i] == NULL) ? NULL : streams[i]->stream; CHECK(ncclReduce(inputs[i].data_ptr(), outputs[i].data_ptr(), count, data_type, (ncclRedOp_t) op, root, comm[i], stream)); } #if defined(NCCL_MAJOR) && (NCCL_MAJOR >= 2) CHECK(ncclGroupEnd()); #endif Py_RETURN_NONE; END_HANDLE_TH_ERRORS }