304 lines
11 KiB
C
304 lines
11 KiB
C
|
#pragma once
|
||
|
|
|
||
|
|
#include <c10/core/DeviceType.h>
|
||
|
|
#include <c10/core/impl/InlineDeviceGuard.h>
|
||
|
|
#include <c10/core/impl/InlineStreamGuard.h>
|
||
|
|
#include <c10/cuda/CUDAMacros.h>
|
||
|
|
#include <c10/cuda/impl/CUDAGuardImpl.h>
|
||
|
|
|
||
|
|
#include <cstddef>
|
||
|
|
|
||
|
|
namespace c10::cuda {
|
||
|
|
|
||
|
|
// This code is kind of boilerplatey. See Note [Whither the DeviceGuard
|
||
|
|
// boilerplate]
|
||
|
|
|
||
|
|
/// A variant of DeviceGuard that is specialized for CUDA. It accepts
|
||
|
|
/// integer indices (interpreting them as CUDA devices) and is a little
|
||
|
|
/// more efficient than DeviceGuard (it compiles to straight line
|
||
|
|
/// cudaSetDevice/cudaGetDevice calls); however, it can only be used
|
||
|
|
/// from code that links against CUDA directly.
|
||
|
|
struct CUDAGuard {
|
||
|
|
/// No default constructor; see Note [Omitted default constructor from RAII]
|
||
|
|
explicit CUDAGuard() = delete;
|
||
|
|
|
||
|
|
/// Set the current CUDA device to the passed device index.
|
||
|
|
explicit CUDAGuard(DeviceIndex device_index) : guard_(device_index) {}
|
||
|
|
|
||
|
|
/// Sets the current CUDA device to the passed device. Errors if the passed
|
||
|
|
/// device is not a CUDA device.
|
||
|
|
explicit CUDAGuard(Device device) : guard_(device) {}
|
||
|
|
|
||
|
|
// Copy is not allowed
|
||
|
|
CUDAGuard(const CUDAGuard&) = delete;
|
||
|
|
CUDAGuard& operator=(const CUDAGuard&) = delete;
|
||
|
|
|
||
|
|
// Move is not allowed (there is no uninitialized state)
|
||
|
|
CUDAGuard(CUDAGuard&& other) = delete;
|
||
|
|
CUDAGuard& operator=(CUDAGuard&& other) = delete;
|
||
|
|
|
||
|
|
/// Sets the CUDA device to the given device. Errors if the given device
|
||
|
|
/// is not a CUDA device.
|
||
|
|
void set_device(Device device) {
|
||
|
|
guard_.set_device(device);
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Sets the CUDA device to the given device. Errors if the given device
|
||
|
|
/// is not a CUDA device. (This method is provided for uniformity with
|
||
|
|
/// DeviceGuard).
|
||
|
|
void reset_device(Device device) {
|
||
|
|
guard_.reset_device(device);
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Sets the CUDA device to the given device index.
|
||
|
|
void set_index(DeviceIndex device_index) {
|
||
|
|
guard_.set_index(device_index);
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Returns the device that was set upon construction of the guard
|
||
|
|
Device original_device() const {
|
||
|
|
return guard_.original_device();
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Returns the last device that was set via `set_device`, if any, otherwise
|
||
|
|
/// the device passed during construction.
|
||
|
|
Device current_device() const {
|
||
|
|
return guard_.current_device();
|
||
|
|
}
|
||
|
|
|
||
|
|
private:
|
||
|
|
/// The guard for the current device.
|
||
|
|
c10::impl::InlineDeviceGuard<impl::CUDAGuardImpl> guard_;
|
||
|
|
};
|
||
|
|
|
||
|
|
/// A variant of OptionalDeviceGuard that is specialized for CUDA. See
|
||
|
|
/// CUDAGuard for when you can use this.
|
||
|
|
struct OptionalCUDAGuard {
|
||
|
|
/// Create an uninitialized OptionalCUDAGuard.
|
||
|
|
explicit OptionalCUDAGuard() : guard_() {}
|
||
|
|
|
||
|
|
/// Set the current CUDA device to the passed Device, if it is not nullopt.
|
||
|
|
explicit OptionalCUDAGuard(optional<Device> device_opt)
|
||
|
|
: guard_(device_opt) {}
|
||
|
|
|
||
|
|
/// Set the current CUDA device to the passed device index, if it is not
|
||
|
|
/// nullopt
|
||
|
|
explicit OptionalCUDAGuard(optional<DeviceIndex> device_index_opt)
|
||
|
|
: guard_(device_index_opt) {}
|
||
|
|
|
||
|
|
// Copy is not allowed
|
||
|
|
OptionalCUDAGuard(const OptionalCUDAGuard&) = delete;
|
||
|
|
OptionalCUDAGuard& operator=(const OptionalCUDAGuard&) = delete;
|
||
|
|
|
||
|
|
// See Note [Move construction for RAII guards is tricky]
|
||
|
|
OptionalCUDAGuard(OptionalCUDAGuard&& other) = delete;
|
||
|
|
|
||
|
|
// See Note [Move assignment for RAII guards is tricky]
|
||
|
|
OptionalCUDAGuard& operator=(OptionalCUDAGuard&& other) = delete;
|
||
|
|
|
||
|
|
/// Sets the CUDA device to the given device, initializing the guard if it
|
||
|
|
/// is not already initialized. Errors if the given device is not a CUDA
|
||
|
|
/// device.
|
||
|
|
void set_device(Device device) {
|
||
|
|
guard_.set_device(device);
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Sets the CUDA device to the given device, initializing the guard if it is
|
||
|
|
/// not already initialized. Errors if the given device is not a CUDA device.
|
||
|
|
/// (This method is provided for uniformity with OptionalDeviceGuard).
|
||
|
|
void reset_device(Device device) {
|
||
|
|
guard_.reset_device(device);
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Sets the CUDA device to the given device index, initializing the guard if
|
||
|
|
/// it is not already initialized.
|
||
|
|
void set_index(DeviceIndex device_index) {
|
||
|
|
guard_.set_index(device_index);
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Returns the device that was set immediately prior to initialization of the
|
||
|
|
/// guard, or nullopt if the guard is uninitialized.
|
||
|
|
optional<Device> original_device() const {
|
||
|
|
return guard_.original_device();
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Returns the most recent device that was set using this device guard,
|
||
|
|
/// either from construction, or via set_device, if the guard is initialized,
|
||
|
|
/// or nullopt if the guard is uninitialized.
|
||
|
|
optional<Device> current_device() const {
|
||
|
|
return guard_.current_device();
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Restore the original CUDA device, resetting this guard to uninitialized
|
||
|
|
/// state.
|
||
|
|
void reset() {
|
||
|
|
guard_.reset();
|
||
|
|
}
|
||
|
|
|
||
|
|
private:
|
||
|
|
c10::impl::InlineOptionalDeviceGuard<impl::CUDAGuardImpl> guard_;
|
||
|
|
};
|
||
|
|
|
||
|
|
/// A variant of StreamGuard that is specialized for CUDA. See CUDAGuard
|
||
|
|
/// for when you can use this.
|
||
|
|
struct CUDAStreamGuard {
|
||
|
|
/// No default constructor, see Note [Omitted default constructor from RAII]
|
||
|
|
explicit CUDAStreamGuard() = delete;
|
||
|
|
|
||
|
|
/// Set the current CUDA device to the device associated with the passed
|
||
|
|
/// stream, and set the current CUDA stream on that device to the passed
|
||
|
|
/// stream. Errors if the Stream is not a CUDA stream.
|
||
|
|
explicit CUDAStreamGuard(Stream stream) : guard_(stream) {}
|
||
|
|
|
||
|
|
/// Copy is disallowed
|
||
|
|
CUDAStreamGuard(const CUDAStreamGuard&) = delete;
|
||
|
|
CUDAStreamGuard& operator=(const CUDAStreamGuard&) = delete;
|
||
|
|
|
||
|
|
/// Move is disallowed, as CUDAStreamGuard does not have an uninitialized
|
||
|
|
/// state, which is required for moves on types with nontrivial destructors.
|
||
|
|
CUDAStreamGuard(CUDAStreamGuard&& other) = delete;
|
||
|
|
CUDAStreamGuard& operator=(CUDAStreamGuard&& other) = delete;
|
||
|
|
|
||
|
|
/// Resets the currently set stream to the original stream and
|
||
|
|
/// the currently set device to the original device. Then,
|
||
|
|
/// set the current device to the device associated with the passed stream,
|
||
|
|
/// and set the current stream on that device to the passed stream.
|
||
|
|
/// Errors if the stream passed is not a CUDA stream.
|
||
|
|
///
|
||
|
|
/// NOTE: this implementation may skip some stream/device setting if
|
||
|
|
/// it can prove that it is unnecessary.
|
||
|
|
///
|
||
|
|
/// WARNING: reset_stream does NOT preserve previously set streams on
|
||
|
|
/// different devices. If you need to set streams on multiple devices
|
||
|
|
/// on CUDA, use CUDAMultiStreamGuard instead.
|
||
|
|
void reset_stream(Stream stream) {
|
||
|
|
guard_.reset_stream(stream);
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Returns the CUDA stream that was set at the time the guard was
|
||
|
|
/// constructed.
|
||
|
|
CUDAStream original_stream() const {
|
||
|
|
return CUDAStream(CUDAStream::UNCHECKED, guard_.original_stream());
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Returns the most recent CUDA stream that was set using this device guard,
|
||
|
|
/// either from construction, or via set_stream.
|
||
|
|
CUDAStream current_stream() const {
|
||
|
|
return CUDAStream(CUDAStream::UNCHECKED, guard_.current_stream());
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Returns the most recent CUDA device that was set using this device guard,
|
||
|
|
/// either from construction, or via set_device/reset_device/set_index.
|
||
|
|
Device current_device() const {
|
||
|
|
return guard_.current_device();
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Returns the CUDA device that was set at the most recent reset_stream(),
|
||
|
|
/// or otherwise the device at construction time.
|
||
|
|
Device original_device() const {
|
||
|
|
return guard_.original_device();
|
||
|
|
}
|
||
|
|
|
||
|
|
private:
|
||
|
|
c10::impl::InlineStreamGuard<impl::CUDAGuardImpl> guard_;
|
||
|
|
};
|
||
|
|
|
||
|
|
/// A variant of OptionalStreamGuard that is specialized for CUDA. See
|
||
|
|
/// CUDAGuard for when you can use this.
|
||
|
|
struct OptionalCUDAStreamGuard {
|
||
|
|
/// Create an uninitialized guard.
|
||
|
|
explicit OptionalCUDAStreamGuard() : guard_() {}
|
||
|
|
|
||
|
|
/// Set the current CUDA device to the device associated with the passed
|
||
|
|
/// stream, and set the current CUDA stream on that device to the passed
|
||
|
|
/// stream. Errors if the Stream is not a CUDA stream.
|
||
|
|
explicit OptionalCUDAStreamGuard(Stream stream) : guard_(stream) {}
|
||
|
|
|
||
|
|
/// Set the current device to the device associated with the passed stream,
|
||
|
|
/// and set the current stream on that device to the passed stream,
|
||
|
|
/// if the passed stream is not nullopt.
|
||
|
|
explicit OptionalCUDAStreamGuard(optional<Stream> stream_opt)
|
||
|
|
: guard_(stream_opt) {}
|
||
|
|
|
||
|
|
/// Copy is disallowed
|
||
|
|
OptionalCUDAStreamGuard(const OptionalCUDAStreamGuard&) = delete;
|
||
|
|
OptionalCUDAStreamGuard& operator=(const OptionalCUDAStreamGuard&) = delete;
|
||
|
|
|
||
|
|
// See Note [Move construction for RAII guards is tricky]
|
||
|
|
OptionalCUDAStreamGuard(OptionalCUDAStreamGuard&& other) = delete;
|
||
|
|
|
||
|
|
// See Note [Move assignment for RAII guards is tricky]
|
||
|
|
OptionalCUDAStreamGuard& operator=(OptionalCUDAStreamGuard&& other) = delete;
|
||
|
|
|
||
|
|
/// Resets the currently set CUDA stream to the original stream and
|
||
|
|
/// the currently set device to the original device. Then,
|
||
|
|
/// set the current device to the device associated with the passed stream,
|
||
|
|
/// and set the current stream on that device to the passed stream.
|
||
|
|
/// Initializes the guard if it was not previously initialized.
|
||
|
|
void reset_stream(Stream stream) {
|
||
|
|
guard_.reset_stream(stream);
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Returns the CUDA stream that was set at the time the guard was most
|
||
|
|
/// recently initialized, or nullopt if the guard is uninitialized.
|
||
|
|
optional<CUDAStream> original_stream() const {
|
||
|
|
auto r = guard_.original_stream();
|
||
|
|
if (r.has_value()) {
|
||
|
|
return make_optional(CUDAStream(CUDAStream::UNCHECKED, r.value()));
|
||
|
|
} else {
|
||
|
|
return nullopt;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Returns the most recent CUDA stream that was set using this stream guard,
|
||
|
|
/// either from construction, or via reset_stream, if the guard is
|
||
|
|
/// initialized, or nullopt if the guard is uninitialized.
|
||
|
|
optional<CUDAStream> current_stream() const {
|
||
|
|
auto r = guard_.current_stream();
|
||
|
|
if (r.has_value()) {
|
||
|
|
return make_optional(CUDAStream(CUDAStream::UNCHECKED, r.value()));
|
||
|
|
} else {
|
||
|
|
return nullopt;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Restore the original CUDA device and stream, resetting this guard to
|
||
|
|
/// uninitialized state.
|
||
|
|
void reset() {
|
||
|
|
guard_.reset();
|
||
|
|
}
|
||
|
|
|
||
|
|
private:
|
||
|
|
c10::impl::InlineOptionalStreamGuard<impl::CUDAGuardImpl> guard_;
|
||
|
|
};
|
||
|
|
|
||
|
|
/// A variant of MultiStreamGuard that is specialized for CUDA.
|
||
|
|
struct CUDAMultiStreamGuard {
|
||
|
|
explicit CUDAMultiStreamGuard(ArrayRef<CUDAStream> streams)
|
||
|
|
: guard_(unwrapStreams(streams)) {}
|
||
|
|
|
||
|
|
/// Copy is disallowed
|
||
|
|
CUDAMultiStreamGuard(const CUDAMultiStreamGuard&) = delete;
|
||
|
|
CUDAMultiStreamGuard& operator=(const CUDAMultiStreamGuard&) = delete;
|
||
|
|
|
||
|
|
// See Note [Move construction for RAII guards is tricky]
|
||
|
|
CUDAMultiStreamGuard(CUDAMultiStreamGuard&& other) = delete;
|
||
|
|
|
||
|
|
// See Note [Move assignment for RAII guards is tricky]
|
||
|
|
CUDAMultiStreamGuard& operator=(CUDAMultiStreamGuard&& other) = delete;
|
||
|
|
|
||
|
|
private:
|
||
|
|
c10::impl::InlineMultiStreamGuard<impl::CUDAGuardImpl> guard_;
|
||
|
|
|
||
|
|
static std::vector<Stream> unwrapStreams(ArrayRef<CUDAStream> cudaStreams) {
|
||
|
|
std::vector<Stream> streams;
|
||
|
|
streams.reserve(cudaStreams.size());
|
||
|
|
for (const CUDAStream& cudaStream : cudaStreams) {
|
||
|
|
streams.push_back(cudaStream);
|
||
|
|
}
|
||
|
|
return streams;
|
||
|
|
}
|
||
|
|
};
|
||
|
|
|
||
|
|
} // namespace c10::cuda
|