Files
mscclpp/src/cuda_utils.cc
2023-06-25 12:40:12 +08:00

20 lines
637 B
C++

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.
#include <mscclpp/cuda_utils.hpp>
namespace mscclpp {
AvoidCudaGraphCaptureGuard::AvoidCudaGraphCaptureGuard() : mode_(cudaStreamCaptureModeRelaxed) {
MSCCLPP_CUDATHROW(cudaThreadExchangeStreamCaptureMode(&mode_));
}
AvoidCudaGraphCaptureGuard::~AvoidCudaGraphCaptureGuard() { cudaThreadExchangeStreamCaptureMode(&mode_); }
CudaStreamWithFlags::CudaStreamWithFlags(unsigned int flags) {
MSCCLPP_CUDATHROW(cudaStreamCreateWithFlags(&stream_, flags));
}
CudaStreamWithFlags::~CudaStreamWithFlags() { cudaStreamDestroy(stream_); }
} // namespace mscclpp