wehub-resource-sync
242 lines
8.8 KiB
C++
242 lines
8.8 KiB
C++
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#pragma once
|
|
|
|
#include "paddle/common/ddim.h"
|
|
#include "paddle/phi/backends/all_context.h"
|
|
#include "paddle/phi/common/memory_utils.h"
|
|
#include "paddle/phi/core/dense_tensor.h"
|
|
#include "paddle/phi/core/tensor_array.h"
|
|
#ifdef PADDLE_WITH_CUDA
|
|
#include "paddle/phi/backends/gpu/cuda/cuda_graph_with_memory_pool.h"
|
|
#endif
|
|
|
|
namespace phi {
|
|
|
|
#if (defined(PADDLE_WITH_RCCL) || defined(PADDLE_WITH_NCCL)) && \
|
|
NCCL_VERSION_CODE >= 2703 || \
|
|
defined(PADDLE_WITH_XPU_BKCL)
|
|
template <typename Context, typename CommContext, typename StreamType>
|
|
void send_shape_info(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
CommContext* comm_ctx,
|
|
int peer,
|
|
StreamType stream) {
|
|
#if (defined(PADDLE_WITH_RCCL) || defined(PADDLE_WITH_NCCL)) && \
|
|
NCCL_VERSION_CODE >= 2703
|
|
PADDLE_ENFORCE_EQ((stream != nullptr && comm_ctx != nullptr),
|
|
true,
|
|
errors::InvalidArgument(
|
|
"NCCLComm and Stream should be provided if use NCCL "
|
|
"to send the shape info."));
|
|
#elif defined(PADDLE_WITH_XPU_BKCL)
|
|
PADDLE_ENFORCE_EQ(
|
|
(comm_ctx != nullptr),
|
|
true,
|
|
errors::InvalidArgument("BKCLComm should be provided if use BKCL "
|
|
"to send the shape info."));
|
|
#endif
|
|
DataType shape_dtype = DataType::INT32;
|
|
auto dims = x.dims();
|
|
int shape_size = dims.size();
|
|
|
|
// step1: send the shape size
|
|
DenseTensor cpu_shape_size_tensor(shape_dtype);
|
|
cpu_shape_size_tensor.Resize({1});
|
|
dev_ctx.HostAlloc(&cpu_shape_size_tensor, shape_dtype);
|
|
auto* cpu_data = cpu_shape_size_tensor.data<int>();
|
|
cpu_data[0] = shape_size;
|
|
|
|
// copy the shape size tensor to gpu/xpu and send
|
|
DenseTensor shape_size_tensor;
|
|
shape_size_tensor.Resize({1});
|
|
dev_ctx.Alloc(&shape_size_tensor, shape_dtype);
|
|
const auto& cpu_place = CPUPlace();
|
|
#ifdef PADDLE_WITH_CUDA
|
|
const int* stable_shape_size =
|
|
phi::backends::gpu::RestoreHostMemIfCapturingCUDAGraph(
|
|
const_cast<int*>(cpu_shape_size_tensor.data<int>()),
|
|
static_cast<size_t>(cpu_shape_size_tensor.numel()));
|
|
memory_utils::Copy(dev_ctx.GetPlace(),
|
|
shape_size_tensor.data(),
|
|
cpu_place,
|
|
stable_shape_size,
|
|
cpu_shape_size_tensor.numel() * sizeof(int),
|
|
stream);
|
|
#else
|
|
memory_utils::Copy(dev_ctx.GetPlace(),
|
|
shape_size_tensor.data(),
|
|
cpu_place,
|
|
cpu_shape_size_tensor.data(),
|
|
cpu_shape_size_tensor.numel() * sizeof(int),
|
|
stream);
|
|
#endif
|
|
|
|
comm_ctx->Send(shape_size_tensor, shape_size_tensor.numel(), peer, stream);
|
|
|
|
// step2: send the shape
|
|
DenseTensor cpu_shape_tensor(shape_dtype);
|
|
cpu_shape_tensor.Resize({shape_size});
|
|
dev_ctx.HostAlloc(&cpu_shape_tensor, shape_dtype);
|
|
auto* cpu_shape_data = cpu_shape_tensor.data<int>();
|
|
for (int i = 0; i < shape_size; ++i) {
|
|
cpu_shape_data[i] = dims[i];
|
|
}
|
|
|
|
// copy the shape tensor to gpu and send
|
|
DenseTensor shape_tensor;
|
|
shape_tensor.Resize({shape_size});
|
|
dev_ctx.Alloc(&shape_tensor, shape_dtype);
|
|
#ifdef PADDLE_WITH_CUDA
|
|
const int* stable_shape =
|
|
phi::backends::gpu::RestoreHostMemIfCapturingCUDAGraph(
|
|
const_cast<int*>(cpu_shape_tensor.data<int>()),
|
|
static_cast<size_t>(cpu_shape_tensor.numel()));
|
|
memory_utils::Copy(dev_ctx.GetPlace(),
|
|
shape_tensor.data(),
|
|
cpu_place,
|
|
stable_shape,
|
|
cpu_shape_tensor.numel() * sizeof(int),
|
|
stream);
|
|
#else
|
|
memory_utils::Copy(dev_ctx.GetPlace(),
|
|
shape_tensor.data(),
|
|
cpu_place,
|
|
cpu_shape_tensor.data(),
|
|
cpu_shape_tensor.numel() * sizeof(int),
|
|
stream);
|
|
#endif
|
|
comm_ctx->Send(shape_tensor, shape_tensor.numel(), peer, stream);
|
|
dev_ctx.Wait();
|
|
}
|
|
#endif
|
|
|
|
#if (defined(PADDLE_WITH_RCCL) || defined(PADDLE_WITH_NCCL)) && \
|
|
NCCL_VERSION_CODE >= 2703 || \
|
|
defined(PADDLE_WITH_XPU_BKCL)
|
|
template <typename Context, typename CommContext, typename StreamType>
|
|
DDim recv_shape_info(const Context& dev_ctx,
|
|
DenseTensor* out,
|
|
CommContext* comm_ctx,
|
|
int peer) {
|
|
StreamType stream = dev_ctx.stream();
|
|
#if (defined(PADDLE_WITH_RCCL) || defined(PADDLE_WITH_NCCL)) && \
|
|
NCCL_VERSION_CODE >= 2703
|
|
PADDLE_ENFORCE_EQ((stream != nullptr && comm_ctx != nullptr),
|
|
true,
|
|
errors::InvalidArgument(
|
|
"NCCLComm and Stream should be provided if use NCCL "
|
|
"to send the shape info."));
|
|
#elif defined(PADDLE_WITH_XPU_BKCL)
|
|
PADDLE_ENFORCE_EQ(
|
|
(comm_ctx != nullptr),
|
|
true,
|
|
errors::InvalidArgument("BKCLComm should be provided if use BKCL "
|
|
"to send the shape info."));
|
|
#endif
|
|
DataType shape_dtype = DataType::INT32;
|
|
|
|
// DenseTensor shape_size_tensortensor(shape_dtype);
|
|
DenseTensor shape_size_tensortensor(shape_dtype);
|
|
shape_size_tensortensor.Resize({1});
|
|
dev_ctx.Alloc(&shape_size_tensortensor, shape_dtype);
|
|
comm_ctx->Recv(
|
|
&shape_size_tensortensor, shape_size_tensortensor.numel(), peer, stream);
|
|
|
|
// copy the shape size tensor to cpu
|
|
#ifdef PADDLE_WITH_CUDA
|
|
PADDLE_ENFORCE_EQ(
|
|
phi::backends::gpu::IsCUDAGraphCapturing(),
|
|
false,
|
|
common::errors::InvalidArgument(
|
|
"RecvShape does not support CUDA Graph capture: async D2H copy to "
|
|
"a locally allocated CPU DenseTensor 'cpu_shape_size_tensor' will "
|
|
"bake the destination address into the graph; on replay the tensor "
|
|
"is re-created at a different address, causing a dangling-pointer "
|
|
"write."));
|
|
#endif
|
|
DenseTensor cpu_shape_size_tensor(shape_dtype);
|
|
cpu_shape_size_tensor.Resize({1});
|
|
dev_ctx.HostAlloc(&cpu_shape_size_tensor, shape_dtype);
|
|
|
|
memory_utils::Copy(CPUPlace(),
|
|
cpu_shape_size_tensor.data(),
|
|
dev_ctx.GetPlace(),
|
|
shape_size_tensortensor.data(),
|
|
shape_size_tensortensor.numel() * sizeof(int),
|
|
stream);
|
|
|
|
auto* cpu_data = cpu_shape_size_tensor.data<int>();
|
|
int shape_size = cpu_data[0];
|
|
|
|
// step2: send the shape
|
|
// DenseTensor shape_tensor(shape_dtype);
|
|
DenseTensor shape_tensor(shape_dtype);
|
|
shape_tensor.Resize({shape_size});
|
|
dev_ctx.Alloc(&shape_tensor, shape_dtype);
|
|
comm_ctx->Recv(&shape_tensor, shape_tensor.numel(), peer, stream);
|
|
|
|
// copy the shape tensor to cpu
|
|
DenseTensor cpu_shape_tensor(shape_dtype);
|
|
cpu_shape_tensor.Resize({shape_size});
|
|
dev_ctx.HostAlloc(&cpu_shape_tensor, shape_dtype);
|
|
|
|
memory_utils::Copy(CPUPlace(),
|
|
cpu_shape_tensor.data(),
|
|
dev_ctx.GetPlace(),
|
|
shape_tensor.data(),
|
|
shape_tensor.numel() * sizeof(int),
|
|
stream);
|
|
dev_ctx.Wait();
|
|
auto* cpu_shape_data = cpu_shape_tensor.data<int>();
|
|
std::vector<int> all_shape;
|
|
for (int i = 0; i < shape_size; ++i) {
|
|
all_shape.emplace_back(cpu_shape_data[i]);
|
|
}
|
|
DDim new_dim;
|
|
new_dim = new_dim.reshape(all_shape);
|
|
|
|
return new_dim;
|
|
}
|
|
|
|
template <typename Context, typename CommContext>
|
|
CommContext* GetCommContext(const Context& dev_ctx, int peer) {
|
|
PADDLE_ENFORCE_GE(
|
|
peer,
|
|
0,
|
|
errors::InvalidArgument("The peer (%d) for send op must be non-negative.",
|
|
peer));
|
|
|
|
auto comm_ctx = static_cast<CommContext*>(dev_ctx.GetCommContext());
|
|
PADDLE_ENFORCE_NE(
|
|
comm_ctx,
|
|
nullptr,
|
|
errors::Unavailable(
|
|
"NCCLCommContext/BKCLCommContext is nullptr, collective op should "
|
|
"has ring_id attr."));
|
|
|
|
PADDLE_ENFORCE_LT(
|
|
peer,
|
|
comm_ctx->GetSize(),
|
|
errors::InvalidArgument("The value of peer (%d) you set must "
|
|
"be less than comm->nranks (%d).",
|
|
peer,
|
|
comm_ctx->GetSize()));
|
|
return comm_ctx;
|
|
}
|
|
#endif
|
|
|
|
} // namespace phi
|