mirror of
https://github.com/microsoft/mscclpp.git
synced 2026-05-25 07:14:40 +00:00
108 lines
3.9 KiB
C++
108 lines
3.9 KiB
C++
// Copyright (c) Microsoft Corporation.
|
|
// Licensed under the MIT license.
|
|
|
|
#include "context.hpp"
|
|
|
|
#include <mscclpp/env.hpp>
|
|
|
|
#include "api.h"
|
|
#include "connection.hpp"
|
|
#include "endpoint.hpp"
|
|
#include "logger.hpp"
|
|
#include "registered_memory.hpp"
|
|
|
|
namespace mscclpp {
|
|
|
|
CudaIpcStream::CudaIpcStream(int deviceId)
|
|
: stream_(std::make_shared<CudaStreamWithFlags>()), deviceId_(deviceId), dirty_(false) {}
|
|
|
|
void CudaIpcStream::setStreamIfNeeded() {
|
|
if (!env()->cudaIpcUseDefaultStream && stream_->empty()) {
|
|
stream_->set(cudaStreamNonBlocking);
|
|
}
|
|
}
|
|
|
|
void CudaIpcStream::memcpyD2D(void* dst, const void* src, size_t nbytes) {
|
|
CudaDeviceGuard deviceGuard(deviceId_);
|
|
setStreamIfNeeded();
|
|
MSCCLPP_CUDATHROW(cudaMemcpyAsync(dst, src, nbytes, cudaMemcpyDeviceToDevice, *stream_));
|
|
dirty_ = true;
|
|
}
|
|
|
|
void CudaIpcStream::memcpyH2D(void* dst, const void* src, size_t nbytes) {
|
|
CudaDeviceGuard deviceGuard(deviceId_);
|
|
setStreamIfNeeded();
|
|
MSCCLPP_CUDATHROW(cudaMemcpyAsync(dst, src, nbytes, cudaMemcpyHostToDevice, *stream_));
|
|
dirty_ = true;
|
|
}
|
|
|
|
void CudaIpcStream::sync() {
|
|
if (dirty_) {
|
|
CudaDeviceGuard deviceGuard(deviceId_);
|
|
setStreamIfNeeded();
|
|
MSCCLPP_CUDATHROW(cudaStreamSynchronize(*stream_));
|
|
dirty_ = false;
|
|
}
|
|
}
|
|
|
|
IbCtx* Context::Impl::getIbContext(Transport ibTransport) {
|
|
// Find IB context or create it
|
|
auto it = ibContexts_.find(ibTransport);
|
|
if (it == ibContexts_.end()) {
|
|
auto ibDev = getIBDeviceName(ibTransport);
|
|
ibContexts_[ibTransport] = std::make_unique<IbCtx>(ibDev);
|
|
return ibContexts_[ibTransport].get();
|
|
}
|
|
return it->second.get();
|
|
}
|
|
|
|
std::shared_ptr<uint64_t> Context::Impl::getToken() {
|
|
if (!tokenPool_) {
|
|
tokenPool_ = std::make_shared<TokenPool>(maxNumTokens_);
|
|
}
|
|
return tokenPool_->getToken();
|
|
}
|
|
|
|
MSCCLPP_API_CPP Context::Context() : pimpl_(std::make_unique<Impl>()) {}
|
|
|
|
MSCCLPP_API_CPP Context::~Context() = default;
|
|
|
|
MSCCLPP_API_CPP RegisteredMemory Context::registerMemory(void* ptr, size_t size, TransportFlags transports) {
|
|
return RegisteredMemory(std::make_shared<RegisteredMemory::Impl>(ptr, size, transports, *pimpl_));
|
|
}
|
|
|
|
MSCCLPP_API_CPP Endpoint Context::createEndpoint(EndpointConfig config) {
|
|
return Endpoint(std::make_shared<Endpoint::Impl>(config, *pimpl_));
|
|
}
|
|
|
|
MSCCLPP_API_CPP Connection Context::connect(const Endpoint& localEndpoint, const Endpoint& remoteEndpoint) {
|
|
if (localEndpoint.device().type == DeviceType::GPU && localEndpoint.device().id < 0) {
|
|
THROW(CONN, Error, ErrorCode::InvalidUsage, "No GPU device ID provided for local endpoint");
|
|
}
|
|
if (remoteEndpoint.device().type == DeviceType::GPU && remoteEndpoint.device().id < 0) {
|
|
THROW(CONN, Error, ErrorCode::InvalidUsage, "No GPU device ID provided for remote endpoint");
|
|
}
|
|
auto localTransport = localEndpoint.transport();
|
|
auto remoteTransport = remoteEndpoint.transport();
|
|
if (localTransport != remoteTransport &&
|
|
!(AllIBTransports.has(localTransport) && AllIBTransports.has(remoteTransport))) {
|
|
THROW(CONN, Error, ErrorCode::InvalidUsage, "Transport mismatch between local (", localTransport, ") and remote (",
|
|
remoteTransport, ") endpoints");
|
|
}
|
|
std::shared_ptr<BaseConnection> conn;
|
|
if (localTransport == Transport::CudaIpc) {
|
|
conn = std::make_shared<CudaIpcConnection>(shared_from_this(), localEndpoint, remoteEndpoint);
|
|
} else if (AllIBTransports.has(localTransport)) {
|
|
conn = std::make_shared<IBConnection>(shared_from_this(), localEndpoint, remoteEndpoint);
|
|
} else if (localTransport == Transport::Ethernet) {
|
|
conn = std::make_shared<EthernetConnection>(shared_from_this(), localEndpoint, remoteEndpoint);
|
|
} else {
|
|
THROW(CONN, Error, ErrorCode::InternalError, "Unsupported transport: ", localTransport,
|
|
" (this usually means EndpointConfig.transport was left at Transport::Unknown — "
|
|
"set it explicitly to CudaIpc, an IB transport, or Ethernet)");
|
|
}
|
|
return Connection(conn);
|
|
}
|
|
|
|
} // namespace mscclpp
|