diff --git a/infini_train/include/common/cuda/kernel_helper.cuh b/infini_train/include/common/cuda/kernel_helper.cuh index 6b532afb..2783f73c 100644 --- a/infini_train/include/common/cuda/kernel_helper.cuh +++ b/infini_train/include/common/cuda/kernel_helper.cuh @@ -109,8 +109,10 @@ template __device__ __forceinline__ T Cos(const T &x) { } template __device__ __forceinline__ T Tanh(const T &x) { - if constexpr (std::is_same_v || std::is_same_v) { - return htanh(x); + if constexpr (std::is_same_v) { + return __float2bfloat16(tanhf(__bfloat162float(x))); + } else if constexpr (std::is_same_v) { + return __float2half(tanhf(__half2float(x))); } else if constexpr (std::is_same_v) { return tanhf(x); } else { diff --git a/infini_train/include/generator.h b/infini_train/include/generator.h new file mode 100644 index 00000000..a676699d --- /dev/null +++ b/infini_train/include/generator.h @@ -0,0 +1,55 @@ +#pragma once + +#include +#include +#include +#include + +#include "infini_train/include/device.h" + +namespace infini_train { + +class GeneratorImpl { +public: + virtual ~GeneratorImpl() = default; + + virtual void ManualSeed(uint64_t seed) = 0; + virtual uint64_t Seed() = 0; + virtual uint64_t InitialSeed() const = 0; + virtual std::vector GetState() const = 0; + virtual void SetState(const std::vector &state) = 0; + virtual Device GetDevice() const = 0; + virtual std::pair ReserveRandomOffset(uint64_t increment) = 0; + + virtual void FillUniform(std::vector &buffer, float from, float to) = 0; + virtual void FillNormal(std::vector &buffer, float mean, float std) = 0; +}; + +class Generator { +public: + explicit Generator(std::shared_ptr impl); + + void ManualSeed(uint64_t seed); + uint64_t Seed(); + uint64_t InitialSeed() const; + std::vector GetState() const; + void SetState(const std::vector &state); + Device GetDevice() const; + std::pair ReserveRandomOffset(uint64_t increment); + + void FillUniform(std::vector &buffer, float from, float to); + void FillNormal(std::vector &buffer, float mean, float std); + +private: + std::shared_ptr impl_; +}; + +std::shared_ptr MakeCPUGenerator(uint64_t seed = 42); +std::shared_ptr MakeCUDAGenerator(int8_t device_index, uint64_t seed = 42); +std::shared_ptr GetDefaultCPUGenerator(); +std::shared_ptr GetDefaultCUDAGenerator(int8_t device_index); +std::shared_ptr GetDefaultGenerator(const Device &device); +void ManualSeed(uint64_t seed); +void ManualSeedAll(uint64_t seed); + +} // namespace infini_train diff --git a/infini_train/include/nn/functional.h b/infini_train/include/nn/functional.h index e4354fd1..dddb297c 100644 --- a/infini_train/include/nn/functional.h +++ b/infini_train/include/nn/functional.h @@ -4,9 +4,12 @@ #include #include +#include "infini_train/include/device.h" + namespace infini_train { +class Generator; class Tensor; -} +} // namespace infini_train namespace infini_train::nn::function { @@ -47,6 +50,12 @@ std::shared_ptr Triu(const std::shared_ptr &input, int64_t diago // A tensor of the given shape filled with the scalar value 1. std::shared_ptr Ones(const std::vector size); +std::shared_ptr Rand(const std::vector &size, Device device = Device(), + std::shared_ptr generator = nullptr); + +std::shared_ptr Randn(const std::vector &size, Device device = Device(), + std::shared_ptr generator = nullptr); + // Returns a new tensor with the reciprocal of the elements of input. // // Args: diff --git a/infini_train/include/nn/init.h b/infini_train/include/nn/init.h index fc6effec..763bdfc2 100644 --- a/infini_train/include/nn/init.h +++ b/infini_train/include/nn/init.h @@ -1,12 +1,11 @@ #pragma once #include -#include -#include #include #include "infini_train/include/datatype.h" #include "infini_train/include/device.h" +#include "infini_train/include/generator.h" namespace infini_train { class Tensor; @@ -15,7 +14,7 @@ class Device; namespace infini_train::nn::init { std::shared_ptr Normal(const std::shared_ptr &tensor, float mean = 0.0, float std = 1.0, - std::optional generator = std::nullopt); + std::shared_ptr generator = nullptr); std::pair CalculateFanInAndFanOut(const std::shared_ptr &tensor); @@ -42,10 +41,10 @@ enum class NonLinearityType : int8_t { std::shared_ptr KaimingUniform(const std::shared_ptr &tensor, float a = 0.0f, KaimingMode mode = KaimingMode::kFanIn, NonLinearityType non_linearity = NonLinearityType::kLeakyReLU, - std::optional generator = std::nullopt); + std::shared_ptr generator = nullptr); std::shared_ptr Uniform(const std::shared_ptr &tensor, float a = 0.0f, float b = 1.0f, - std::optional generator = std::nullopt); + std::shared_ptr generator = nullptr); std::shared_ptr Ones(const std::shared_ptr &tensor); diff --git a/infini_train/include/tensor.h b/infini_train/include/tensor.h index dcfd8927..fe066ed3 100644 --- a/infini_train/include/tensor.h +++ b/infini_train/include/tensor.h @@ -4,7 +4,6 @@ #include #include #include -#include #include #include "Eigen/Dense" @@ -15,6 +14,7 @@ #include "infini_train/include/scalar.h" namespace infini_train { +class Generator; namespace autograd { class Function; class AccumulateGrad; @@ -150,8 +150,7 @@ class Tensor : public std::enable_shared_from_this { std::shared_ptr Unsqueeze(int64_t dim); // distribution - std::shared_ptr Uniform(float from = 0.0f, float to = 1.0f, - std::optional generator = std::nullopt); + std::shared_ptr Uniform(float from = 0.0f, float to = 1.0f, std::shared_ptr generator = nullptr); std::shared_ptr Matmul(const std::shared_ptr &other); std::shared_ptr Outer(const std::shared_ptr &other); diff --git a/infini_train/src/generator.cc b/infini_train/src/generator.cc new file mode 100644 index 00000000..6228b0e1 --- /dev/null +++ b/infini_train/src/generator.cc @@ -0,0 +1,433 @@ +#include "infini_train/include/generator.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +namespace infini_train { +namespace { +constexpr char kCPUStateHeader[] = "InfiniTrain CPUGeneratorImpl v1"; +constexpr char kCUDAStateHeader[] = "InfiniTrain CUDAGeneratorImpl v3"; +constexpr float kTwoPi = 6.283185307179586476925286766559f; + +struct Philox4x32State { + uint32_t c0; + uint32_t c1; + uint32_t c2; + uint32_t c3; +}; + +uint32_t MulHi(uint32_t a, uint32_t b) { return static_cast((static_cast(a) * b) >> 32); } + +Philox4x32State PhiloxRound(Philox4x32State counter, uint32_t key0, uint32_t key1) { + constexpr uint32_t kPhiloxM0 = 0xD2511F53; + constexpr uint32_t kPhiloxM1 = 0xCD9E8D57; + const uint32_t lo0 = counter.c0 * kPhiloxM0; + const uint32_t hi0 = MulHi(counter.c0, kPhiloxM0); + const uint32_t lo1 = counter.c2 * kPhiloxM1; + const uint32_t hi1 = MulHi(counter.c2, kPhiloxM1); + return {hi1 ^ counter.c1 ^ key0, lo1, hi0 ^ counter.c3 ^ key1, lo0}; +} + +Philox4x32State Philox(uint64_t seed, uint64_t subsequence) { + constexpr uint32_t kPhiloxW0 = 0x9E3779B9; + constexpr uint32_t kPhiloxW1 = 0xBB67AE85; + Philox4x32State counter{static_cast(subsequence), static_cast(subsequence >> 32), 0, 0}; + uint32_t key0 = static_cast(seed); + uint32_t key1 = static_cast(seed >> 32); + for (int round = 0; round < 10; ++round) { + counter = PhiloxRound(counter, key0, key1); + key0 += kPhiloxW0; + key1 += kPhiloxW1; + } + return counter; +} + +uint32_t PhiloxRandomUint(uint64_t seed, uint64_t offset) { + const auto values = Philox(seed, offset / 4); + switch (offset % 4) { + case 0: + return values.c0; + case 1: + return values.c1; + case 2: + return values.c2; + default: + return values.c3; + } +} + +float UintToUniform(uint32_t value) { return static_cast(value >> 8) * 0x1.0p-24f; } + +void CheckUniformBounds(float from, float to) { + CHECK_LE(from, to); + CHECK(std::isfinite(from)) << "Uniform lower bound must be finite"; + CHECK(std::isfinite(to)) << "Uniform upper bound must be finite"; + const double range = static_cast(to) - static_cast(from); + CHECK_LE(range, static_cast(std::numeric_limits::max())) + << "Uniform bounds range exceeds float maximum"; +} + +void SeedEngine(std::mt19937 &engine, uint64_t seed) { + std::seed_seq seed_seq{static_cast(seed), static_cast(seed >> 32)}; + engine.seed(seed_seq); +} + +uint64_t MakeRandomSeed() { + std::random_device rd; + return (static_cast(rd()) << 32) ^ rd(); +} + +std::string ToString(const std::vector &state) { return std::string(state.begin(), state.end()); } + +std::vector ToBytes(const std::string &state) { return std::vector(state.begin(), state.end()); } + +template T ParseStateIntegerToken(const std::string &token, const char *error_message) { + T value = 0; + const auto [ptr, error] = std::from_chars(token.data(), token.data() + token.size(), value); + CHECK(error == std::errc{} && ptr == token.data() + token.size()) << error_message; + return value; +} + +template T ParseStateIntegerLine(std::istringstream &iss, const char *error_message) { + std::string line; + CHECK(std::getline(iss, line)) << error_message; + return ParseStateIntegerToken(line, error_message); +} + +void ValidateCPUEngineState(const std::string &serialized) { + std::istringstream iss(serialized); + bool any_nonzero_word = false; + for (size_t i = 0; i < std::mt19937::state_size; ++i) { + std::string token; + CHECK(iss >> token) << "Invalid CPU generator engine state: missing state word"; + const uint64_t word = ParseStateIntegerToken(token, "Invalid CPU generator engine word in state"); + CHECK_LE(word, std::numeric_limits::max()) << "Invalid CPU generator engine word in state"; + any_nonzero_word = any_nonzero_word || word != 0; + } + + std::string position_token; + CHECK(iss >> position_token) << "Invalid CPU generator engine state: missing position"; + const uint64_t position + = ParseStateIntegerToken(position_token, "Invalid CPU generator engine position in state"); + CHECK_LE(position, std::mt19937::state_size) << "Invalid CPU generator engine position in state"; + + std::string trailing_token; + CHECK(!(iss >> trailing_token)) << "Invalid CPU generator state: trailing data"; + CHECK(any_nonzero_word) << "Invalid CPU generator engine state: all-zero state"; +} + +void CheckNoTrailingStateData(std::istringstream &iss, const char *generator_name) { + CHECK_EQ(iss.peek(), std::char_traits::eof()) + << "Invalid " << generator_name << " generator state: trailing data"; +} + +std::mutex &DefaultGeneratorMutex() { + static std::mutex mutex; + return mutex; +} + +uint64_t ProcessDefaultSeed() { + static const uint64_t seed = MakeRandomSeed(); + return seed; +} + +uint64_t &DefaultGeneratorSeed() { + static uint64_t seed = ProcessDefaultSeed(); + return seed; +} + +std::unordered_map> &DefaultCUDAGenerators() { + static std::unordered_map> generators; + return generators; +} +} // namespace + +class CPUGeneratorImpl : public GeneratorImpl { +public: + explicit CPUGeneratorImpl(uint64_t seed); + + void ManualSeed(uint64_t seed) override; + uint64_t Seed() override; + uint64_t InitialSeed() const override; + std::vector GetState() const override; + void SetState(const std::vector &state) override; + Device GetDevice() const override; + std::pair ReserveRandomOffset(uint64_t increment) override; + + void FillUniform(std::vector &buffer, float from, float to) override; + void FillNormal(std::vector &buffer, float mean, float std) override; + +private: + mutable std::mutex mutex_; + uint64_t initial_seed_ = 0; + std::mt19937 engine_; +}; + +class CUDAGeneratorImpl : public GeneratorImpl { +public: + CUDAGeneratorImpl(int8_t device_index, uint64_t seed); + + void ManualSeed(uint64_t seed) override; + uint64_t Seed() override; + uint64_t InitialSeed() const override; + std::vector GetState() const override; + void SetState(const std::vector &state) override; + Device GetDevice() const override; + std::pair ReserveRandomOffset(uint64_t increment) override; + + void FillUniform(std::vector &buffer, float from, float to) override; + void FillNormal(std::vector &buffer, float mean, float std) override; + +private: + mutable std::mutex mutex_; + Device device_; + uint64_t initial_seed_ = 0; + uint64_t offset_ = 0; +}; + +Generator::Generator(std::shared_ptr impl) : impl_(std::move(impl)) { CHECK(impl_ != nullptr); } + +void Generator::ManualSeed(uint64_t seed) { impl_->ManualSeed(seed); } + +uint64_t Generator::Seed() { return impl_->Seed(); } + +uint64_t Generator::InitialSeed() const { return impl_->InitialSeed(); } + +std::vector Generator::GetState() const { return impl_->GetState(); } + +void Generator::SetState(const std::vector &state) { impl_->SetState(state); } + +Device Generator::GetDevice() const { return impl_->GetDevice(); } + +std::pair Generator::ReserveRandomOffset(uint64_t increment) { + return impl_->ReserveRandomOffset(increment); +} + +void Generator::FillUniform(std::vector &buffer, float from, float to) { impl_->FillUniform(buffer, from, to); } + +void Generator::FillNormal(std::vector &buffer, float mean, float std) { impl_->FillNormal(buffer, mean, std); } + +CPUGeneratorImpl::CPUGeneratorImpl(uint64_t seed) { ManualSeed(seed); } + +void CPUGeneratorImpl::ManualSeed(uint64_t seed) { + std::lock_guard lock(mutex_); + initial_seed_ = seed; + SeedEngine(engine_, seed); +} + +uint64_t CPUGeneratorImpl::Seed() { + const uint64_t seed = MakeRandomSeed(); + ManualSeed(seed); + return seed; +} + +uint64_t CPUGeneratorImpl::InitialSeed() const { + std::lock_guard lock(mutex_); + return initial_seed_; +} + +std::vector CPUGeneratorImpl::GetState() const { + std::lock_guard lock(mutex_); + std::ostringstream oss; + oss << kCPUStateHeader << "\n" << initial_seed_ << "\n" << engine_; + return ToBytes(oss.str()); +} + +void CPUGeneratorImpl::SetState(const std::vector &state) { + const std::string serialized = ToString(state); + std::istringstream iss(serialized); + + std::string header; + std::mt19937 engine; + std::getline(iss, header); + CHECK_EQ(header, kCPUStateHeader) << "Invalid CPU generator state header"; + const uint64_t seed = ParseStateIntegerLine(iss, "Invalid CPU generator seed in state"); + std::ostringstream engine_state_stream; + engine_state_stream << iss.rdbuf(); + const std::string engine_state = engine_state_stream.str(); + ValidateCPUEngineState(engine_state); + std::istringstream engine_iss(engine_state); + CHECK(engine_iss >> engine) << "Invalid CPU generator engine state"; + CheckNoTrailingStateData(engine_iss, "CPU"); + + std::lock_guard lock(mutex_); + initial_seed_ = seed; + engine_ = engine; +} + +Device CPUGeneratorImpl::GetDevice() const { return Device(); } + +std::pair CPUGeneratorImpl::ReserveRandomOffset(uint64_t) { + LOG(FATAL) << "CPU generator does not expose a CUDA random offset"; + return {0, 0}; +} + +void CPUGeneratorImpl::FillUniform(std::vector &buffer, float from, float to) { + CheckUniformBounds(from, to); + std::lock_guard lock(mutex_); + std::uniform_real_distribution dis(from, to); + std::generate(buffer.begin(), buffer.end(), [&]() { return dis(engine_); }); +} + +void CPUGeneratorImpl::FillNormal(std::vector &buffer, float mean, float std) { + CHECK_GE(std, 0.0f); + std::lock_guard lock(mutex_); + if (std == 0.0f) { + std::normal_distribution dis(0.0f, 1.0f); + std::generate(buffer.begin(), buffer.end(), [&]() { + (void)dis(engine_); + return mean; + }); + return; + } + std::normal_distribution dis(mean, std); + std::generate(buffer.begin(), buffer.end(), [&]() { return dis(engine_); }); +} + +CUDAGeneratorImpl::CUDAGeneratorImpl(int8_t device_index, uint64_t seed) + : device_(Device::DeviceType::kCUDA, device_index) { + CHECK_GE(device_index, 0); + ManualSeed(seed); +} + +void CUDAGeneratorImpl::ManualSeed(uint64_t seed) { + std::lock_guard lock(mutex_); + initial_seed_ = seed; + offset_ = 0; +} + +uint64_t CUDAGeneratorImpl::Seed() { + const uint64_t seed = MakeRandomSeed(); + ManualSeed(seed); + return seed; +} + +uint64_t CUDAGeneratorImpl::InitialSeed() const { + std::lock_guard lock(mutex_); + return initial_seed_; +} + +std::vector CUDAGeneratorImpl::GetState() const { + std::lock_guard lock(mutex_); + std::ostringstream oss; + oss << kCUDAStateHeader << "\n" + << initial_seed_ << "\n" + << static_cast(device_.index()) << "\n" + << offset_ << "\n"; + return ToBytes(oss.str()); +} + +void CUDAGeneratorImpl::SetState(const std::vector &state) { + const std::string serialized = ToString(state); + std::istringstream iss(serialized); + + std::string header; + std::getline(iss, header); + CHECK_EQ(header, kCUDAStateHeader) << "Invalid CUDA generator state header"; + const uint64_t seed = ParseStateIntegerLine(iss, "Invalid CUDA generator seed in state"); + const int device_index = ParseStateIntegerLine(iss, "Invalid CUDA generator device index in state"); + CHECK_GE(device_index, 0) << "Invalid CUDA generator device index in state"; + CHECK_LE(device_index, std::numeric_limits::max()) << "Invalid CUDA generator device index in state"; + const uint64_t offset = ParseStateIntegerLine(iss, "Invalid CUDA generator offset in state"); + CheckNoTrailingStateData(iss, "CUDA"); + + std::lock_guard lock(mutex_); + initial_seed_ = seed; + offset_ = offset; +} + +Device CUDAGeneratorImpl::GetDevice() const { return device_; } + +std::pair CUDAGeneratorImpl::ReserveRandomOffset(uint64_t increment) { + std::lock_guard lock(mutex_); + CHECK_LE(increment, std::numeric_limits::max() - offset_) << "CUDA generator offset overflow"; + const uint64_t offset = offset_; + offset_ += increment; + return {initial_seed_, offset}; +} + +void CUDAGeneratorImpl::FillUniform(std::vector &buffer, float from, float to) { + CheckUniformBounds(from, to); + std::lock_guard lock(mutex_); + CHECK_LE(buffer.size(), std::numeric_limits::max() - offset_) << "CUDA generator offset overflow"; + for (size_t i = 0; i < buffer.size(); ++i) { + const float uniform = UintToUniform(PhiloxRandomUint(initial_seed_, offset_ + i)); + buffer[i] = from + (to - from) * uniform; + } + offset_ += buffer.size(); +} + +void CUDAGeneratorImpl::FillNormal(std::vector &buffer, float mean, float stddev) { + CHECK_GE(stddev, 0.0f); + std::lock_guard lock(mutex_); + CHECK_LE(buffer.size(), std::numeric_limits::max() / 2) << "CUDA generator offset overflow"; + CHECK_LE(buffer.size() * 2, std::numeric_limits::max() - offset_) << "CUDA generator offset overflow"; + for (size_t i = 0; i < buffer.size(); ++i) { + const uint64_t element_offset = offset_ + i * 2; + float uniform1 = UintToUniform(PhiloxRandomUint(initial_seed_, element_offset)); + const float uniform2 = UintToUniform(PhiloxRandomUint(initial_seed_, element_offset + 1)); + uniform1 = std::max(uniform1, 0x1.0p-24f); + const float normal = std::sqrt(-2.0f * std::log(uniform1)) * std::cos(kTwoPi * uniform2); + buffer[i] = mean + stddev * normal; + } + offset_ += buffer.size() * 2; +} + +std::shared_ptr MakeCPUGenerator(uint64_t seed) { + return std::make_shared(std::make_shared(seed)); +} + +std::shared_ptr MakeCUDAGenerator(int8_t device_index, uint64_t seed) { + return std::make_shared(std::make_shared(device_index, seed)); +} + +std::shared_ptr GetDefaultCPUGenerator() { + static auto generator = MakeCPUGenerator(ProcessDefaultSeed()); + return generator; +} + +std::shared_ptr GetDefaultCUDAGenerator(int8_t device_index) { + CHECK_GE(device_index, 0); + std::lock_guard lock(DefaultGeneratorMutex()); + auto &generators = DefaultCUDAGenerators(); + auto it = generators.find(device_index); + if (it != generators.end()) { + return it->second; + } + auto generator = MakeCUDAGenerator(device_index, DefaultGeneratorSeed()); + generators.emplace(device_index, generator); + return generator; +} + +std::shared_ptr GetDefaultGenerator(const Device &device) { + if (device.IsCPU()) { + return GetDefaultCPUGenerator(); + } + if (device.IsCUDA()) { + return GetDefaultCUDAGenerator(device.index()); + } + LOG(FATAL) << "Unsupported default Generator device: " << device; + return nullptr; +} + +void ManualSeed(uint64_t seed) { ManualSeedAll(seed); } + +void ManualSeedAll(uint64_t seed) { + GetDefaultCPUGenerator()->ManualSeed(seed); + + std::lock_guard lock(DefaultGeneratorMutex()); + DefaultGeneratorSeed() = seed; + for (auto &[_, generator] : DefaultCUDAGenerators()) { generator->ManualSeed(seed); } +} + +} // namespace infini_train diff --git a/infini_train/src/kernels/cpu/cross_entropy.cc b/infini_train/src/kernels/cpu/cross_entropy.cc index f520b2e9..5cf95cf2 100644 --- a/infini_train/src/kernels/cpu/cross_entropy.cc +++ b/infini_train/src/kernels/cpu/cross_entropy.cc @@ -2,6 +2,7 @@ #include #include #include +#include #include #include "glog/logging.h" diff --git a/infini_train/src/kernels/cpu/transform.cc b/infini_train/src/kernels/cpu/transform.cc index 1a810b44..04031631 100644 --- a/infini_train/src/kernels/cpu/transform.cc +++ b/infini_train/src/kernels/cpu/transform.cc @@ -1,5 +1,6 @@ #include #include +#include #include "glog/logging.h" diff --git a/infini_train/src/kernels/cuda/elementwise.cu b/infini_train/src/kernels/cuda/elementwise.cu index fc423b35..96f3bda1 100644 --- a/infini_train/src/kernels/cuda/elementwise.cu +++ b/infini_train/src/kernels/cuda/elementwise.cu @@ -1,6 +1,7 @@ #include #include +#include #include "infini_train/include/common/common.h" #include "infini_train/include/common/cuda/common_cuda.h" diff --git a/infini_train/src/kernels/cuda/gather.cu b/infini_train/src/kernels/cuda/gather.cu index d8b0cffa..a4790035 100644 --- a/infini_train/src/kernels/cuda/gather.cu +++ b/infini_train/src/kernels/cuda/gather.cu @@ -1,4 +1,5 @@ #include "glog/logging.h" +#include #include "infini_train/include/common/common.h" #include "infini_train/include/common/cuda/common_cuda.h" diff --git a/infini_train/src/kernels/cuda/no_op.cu b/infini_train/src/kernels/cuda/no_op.cu index ef2c9566..e12fd613 100644 --- a/infini_train/src/kernels/cuda/no_op.cu +++ b/infini_train/src/kernels/cuda/no_op.cu @@ -1,4 +1,5 @@ #include "glog/logging.h" +#include #include "infini_train/include/dispatcher.h" #include "infini_train/include/tensor.h" diff --git a/infini_train/src/kernels/cuda/random.cu b/infini_train/src/kernels/cuda/random.cu new file mode 100644 index 00000000..2b271760 --- /dev/null +++ b/infini_train/src/kernels/cuda/random.cu @@ -0,0 +1,142 @@ +#include +#include +#include + +#include "infini_train/include/common/cuda/common_cuda.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/device.h" + +#include "infini_train/src/core/runtime/cuda/cuda_runtime_common.h" + +namespace infini_train::kernels::cuda { +namespace { +constexpr int kThreadsPerBlock = 256; +constexpr float kTwoPi = 6.283185307179586476925286766559f; + +void CheckUniformBounds(float from, float to) { + CHECK_LE(from, to); + CHECK(std::isfinite(from)) << "Uniform lower bound must be finite"; + CHECK(std::isfinite(to)) << "Uniform upper bound must be finite"; + const double range = static_cast(to) - static_cast(from); + CHECK_LE(range, static_cast(std::numeric_limits::max())) + << "Uniform bounds range exceeds float maximum"; +} + +struct Philox4x32State { + uint32_t c0; + uint32_t c1; + uint32_t c2; + uint32_t c3; +}; + +__device__ uint32_t MulHi(uint32_t a, uint32_t b) { + return static_cast((static_cast(a) * b) >> 32); +} + +__device__ Philox4x32State PhiloxRound(Philox4x32State counter, uint32_t key0, uint32_t key1) { + constexpr uint32_t kPhiloxM0 = 0xD2511F53; + constexpr uint32_t kPhiloxM1 = 0xCD9E8D57; + + const uint32_t lo0 = counter.c0 * kPhiloxM0; + const uint32_t hi0 = MulHi(counter.c0, kPhiloxM0); + const uint32_t lo1 = counter.c2 * kPhiloxM1; + const uint32_t hi1 = MulHi(counter.c2, kPhiloxM1); + + return {hi1 ^ counter.c1 ^ key0, lo1, hi0 ^ counter.c3 ^ key1, lo0}; +} + +__device__ Philox4x32State Philox(uint64_t seed, uint64_t subsequence) { + constexpr uint32_t kPhiloxW0 = 0x9E3779B9; + constexpr uint32_t kPhiloxW1 = 0xBB67AE85; + + Philox4x32State counter{static_cast(subsequence), static_cast(subsequence >> 32), 0, 0}; + uint32_t key0 = static_cast(seed); + uint32_t key1 = static_cast(seed >> 32); + + for (int round = 0; round < 10; ++round) { + counter = PhiloxRound(counter, key0, key1); + key0 += kPhiloxW0; + key1 += kPhiloxW1; + } + return counter; +} + +__device__ uint32_t PhiloxRandomUint(uint64_t seed, uint64_t offset) { + const Philox4x32State values = Philox(seed, offset / 4); + switch (offset % 4) { + case 0: + return values.c0; + case 1: + return values.c1; + case 2: + return values.c2; + default: + return values.c3; + } +} + +__device__ float UintToUniform(uint32_t value) { return static_cast(value >> 8) * 0x1.0p-24f; } + +__global__ void RandomUniformFloat32Kernel(float *data, int64_t num_elements, float from, float to, uint64_t seed, + uint64_t offset) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= num_elements) { + return; + } + const float u = UintToUniform(PhiloxRandomUint(seed, offset + static_cast(idx))); + data[idx] = from + (to - from) * u; +} + +__global__ void RandomNormalFloat32Kernel(float *data, int64_t num_elements, float mean, float std, uint64_t seed, + uint64_t offset) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= num_elements) { + return; + } + + const uint64_t element_offset = offset + static_cast(idx) * 2; + float u1 = UintToUniform(PhiloxRandomUint(seed, element_offset)); + const float u2 = UintToUniform(PhiloxRandomUint(seed, element_offset + 1)); + u1 = fmaxf(u1, 0x1.0p-24f); + const float z = sqrtf(-2.0f * logf(u1)) * cosf(kTwoPi * u2); + data[idx] = mean + std * z; +} + +cudaStream_t GetCudaStream(Device device) { + auto *stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)); + CHECK_NOTNULL(stream); + return stream->cuda_stream(); +} +} // namespace + +void RandomUniformFloat32(void *data, int64_t num_elements, float from, float to, uint64_t seed, uint64_t offset, + Device device) { + CHECK_GE(num_elements, 0); + CheckUniformBounds(from, to); + if (num_elements == 0) { + return; + } + CHECK_LE(num_elements, static_cast(std::numeric_limits::max()) * kThreadsPerBlock) + << "Random uniform tensor is too large"; + const int blocks = static_cast((num_elements + kThreadsPerBlock - 1) / kThreadsPerBlock); + RandomUniformFloat32Kernel<<>>( + static_cast(data), num_elements, from, to, seed, offset); + CUDA_CHECK(cudaGetLastError()); +} + +void RandomNormalFloat32(void *data, int64_t num_elements, float mean, float std, uint64_t seed, uint64_t offset, + Device device) { + CHECK_GE(num_elements, 0); + CHECK_GE(std, 0.0f); + if (num_elements == 0) { + return; + } + CHECK_LE(num_elements, static_cast(std::numeric_limits::max()) * kThreadsPerBlock) + << "Random normal tensor is too large"; + const int blocks = static_cast((num_elements + kThreadsPerBlock - 1) / kThreadsPerBlock); + RandomNormalFloat32Kernel<<>>( + static_cast(data), num_elements, mean, std, seed, offset); + CUDA_CHECK(cudaGetLastError()); +} +} // namespace infini_train::kernels::cuda diff --git a/infini_train/src/kernels/cuda/reduction.cu b/infini_train/src/kernels/cuda/reduction.cu index c56470e3..793e001b 100644 --- a/infini_train/src/kernels/cuda/reduction.cu +++ b/infini_train/src/kernels/cuda/reduction.cu @@ -1,4 +1,5 @@ #include +#include #include "infini_train/include/common/cuda/common_cuda.h" #include "infini_train/include/common/cuda/cub_compat.cuh" diff --git a/infini_train/src/nn/functional.cc b/infini_train/src/nn/functional.cc index c33e2368..2b912a1a 100644 --- a/infini_train/src/nn/functional.cc +++ b/infini_train/src/nn/functional.cc @@ -9,6 +9,7 @@ #include "infini_train/include/autograd/reduction.h" #include "infini_train/include/autograd/softmax.h" #include "infini_train/include/autograd/transform.h" +#include "infini_train/include/generator.h" #include "infini_train/include/nn/init.h" #include "infini_train/include/tensor.h" @@ -26,6 +27,16 @@ std::shared_ptr Ones(const std::vector size) { return init::Ones(ones); } +std::shared_ptr Rand(const std::vector &size, Device device, std::shared_ptr generator) { + auto tensor = std::make_shared(size, DataType::kFLOAT32, device); + return init::Uniform(tensor, 0.0f, 1.0f, generator); +} + +std::shared_ptr Randn(const std::vector &size, Device device, std::shared_ptr generator) { + auto tensor = std::make_shared(size, DataType::kFLOAT32, device); + return init::Normal(tensor, 0.0f, 1.0f, generator); +} + std::shared_ptr Reciprocal(const std::shared_ptr &input) { return input->Reciprocal(); } std::shared_ptr Sin(const std::shared_ptr &input) { return input->Sin(); } diff --git a/infini_train/src/nn/init.cc b/infini_train/src/nn/init.cc index 79b4b48b..a30ac6c2 100644 --- a/infini_train/src/nn/init.cc +++ b/infini_train/src/nn/init.cc @@ -1,61 +1,68 @@ #include "infini_train/include/nn/init.h" -#include +#include #include +#include #include +#include #include -#include #include -#ifdef USE_OMP -#include -#endif - #include "glog/logging.h" #include "infini_train/include/core/runtime/device_guard.h" #include "infini_train/include/device.h" #include "infini_train/include/tensor.h" +#ifdef USE_CUDA +namespace infini_train::kernels::cuda { +void RandomUniformFloat32(void *data, int64_t num_elements, float from, float to, uint64_t seed, uint64_t offset, + Device device); +void RandomNormalFloat32(void *data, int64_t num_elements, float mean, float std, uint64_t seed, uint64_t offset, + Device device); +} // namespace infini_train::kernels::cuda +#endif + namespace infini_train::nn::init { namespace { -constexpr int kRandomSeed = 42; - -// FIXME: RNG design is incomplete. -// -// Current implementation lacks: -// - unified Generator abstraction -// - global default generator and seed control -// - reproducible / clonable RNG state -// -// TODO: -// - introduce Generator interface and backend impl -// - add default generator management (per device) -// - refactor random ops to consume Generator -static std::mt19937 gen(kRandomSeed); +void CheckUniformBounds(float from, float to) { + CHECK_LE(from, to); + CHECK(std::isfinite(from)) << "Uniform lower bound must be finite"; + CHECK(std::isfinite(to)) << "Uniform upper bound must be finite"; + const double range = static_cast(to) - static_cast(from); + CHECK_LE(range, static_cast(std::numeric_limits::max())) + << "Uniform bounds range exceeds float maximum"; +} } // namespace std::shared_ptr Normal(const std::shared_ptr &tensor, float mean, float std, - std::optional generator) { + std::shared_ptr generator) { + CHECK_GE(std, 0.0f); + CHECK(tensor->Dtype() == DataType::kFLOAT32) << "Random normal currently supports float32 tensors"; const int64_t num_elements = tensor->NumElements(); - std::vector buffer(num_elements); -#ifdef USE_OMP -#pragma omp parallel - { - std::mt19937 local_gen(kRandomSeed + omp_get_thread_num()); - std::normal_distribution local_dis(mean, std); -#pragma omp for - for (int i = 0; i < buffer.size(); ++i) { - buffer[i] = generator ? local_dis(generator.value()) : local_dis(local_gen); - } + auto device = tensor->GetDevice(); + auto resolved_generator = generator ? generator : GetDefaultGenerator(device); + CHECK(resolved_generator->GetDevice().type() == device.type()) + << "Generator backend must match tensor device backend: generator=" << resolved_generator->GetDevice() + << " tensor=" << device; + if (num_elements == 0) { + return tensor; + } + +#ifdef USE_CUDA + if (device.IsCUDA()) { + CHECK_LE(num_elements, static_cast(std::numeric_limits::max() / 2)) + << "Random normal tensor is too large"; + const auto [seed, offset] = resolved_generator->ReserveRandomOffset(static_cast(num_elements) * 2); + core::DeviceGuard guard(device); + kernels::cuda::RandomNormalFloat32(tensor->DataPtr(), num_elements, mean, std, seed, offset, device); + return tensor; } -#else - std::normal_distribution dis(mean, std); - std::generate(buffer.begin(), buffer.end(), [&]() { return generator ? dis(generator.value()) : dis(gen); }); #endif - auto device = tensor->GetDevice(); + std::vector buffer(num_elements); + resolved_generator->FillNormal(buffer, mean, std); core::DeviceGuard guard(device); auto impl = core::GetDeviceGuardImpl(device.type()); @@ -113,7 +120,7 @@ float CalculateGain(NonLinearityType nonlinearity, std::optional param = } // namespace std::shared_ptr KaimingUniform(const std::shared_ptr &tensor, float a, KaimingMode mode, - NonLinearityType nonlinearity, std::optional generator) { + NonLinearityType nonlinearity, std::shared_ptr generator) { for (const auto dim : tensor->Dims()) { if (dim == 0) { LOG(WARNING) << "Initializing zero-element tensors is a no-op"; @@ -128,26 +135,31 @@ std::shared_ptr KaimingUniform(const std::shared_ptr &tensor, fl } std::shared_ptr Uniform(const std::shared_ptr &tensor, float a, float b, - std::optional generator) { + std::shared_ptr generator) { + CheckUniformBounds(a, b); + CHECK(tensor->Dtype() == DataType::kFLOAT32) << "Random uniform currently supports float32 tensors"; const int64_t num_elements = tensor->NumElements(); - std::vector buffer(num_elements); -#ifdef USE_OMP -#pragma omp parallel - { - std::mt19937 local_gen(kRandomSeed + omp_get_thread_num()); - std::uniform_real_distribution local_dis(a, b); -#pragma omp for - for (int i = 0; i < buffer.size(); ++i) { - buffer[i] = generator ? local_dis(generator.value()) : local_dis(local_gen); - } + auto device = tensor->GetDevice(); + auto resolved_generator = generator ? generator : GetDefaultGenerator(device); + CHECK(resolved_generator->GetDevice().type() == device.type()) + << "Generator backend must match tensor device backend: generator=" << resolved_generator->GetDevice() + << " tensor=" << device; + if (num_elements == 0) { + return tensor; + } + +#ifdef USE_CUDA + if (device.IsCUDA()) { + const auto [seed, offset] = resolved_generator->ReserveRandomOffset(static_cast(num_elements)); + core::DeviceGuard guard(device); + kernels::cuda::RandomUniformFloat32(tensor->DataPtr(), num_elements, a, b, seed, offset, device); + return tensor; } -#else - std::uniform_real_distribution dis(a, b); - std::generate(buffer.begin(), buffer.end(), [&]() { return generator ? dis(generator.value()) : dis(gen); }); #endif - auto device = tensor->GetDevice(); + std::vector buffer(num_elements); + resolved_generator->FillUniform(buffer, a, b); core::DeviceGuard guard(device); auto impl = core::GetDeviceGuardImpl(device.type()); diff --git a/infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc b/infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc index ab3a8002..b5afd48b 100644 --- a/infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc +++ b/infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc @@ -3,6 +3,7 @@ #include #include #include +#include #include "glog/logging.h" diff --git a/infini_train/src/tensor.cc b/infini_train/src/tensor.cc index 18ca3d22..b454a0e5 100644 --- a/infini_train/src/tensor.cc +++ b/infini_train/src/tensor.cc @@ -471,7 +471,7 @@ std::shared_ptr Tensor::Outer(const std::shared_ptr &other) { } // distribution -std::shared_ptr Tensor::Uniform(float from, float to, std::optional generator) { +std::shared_ptr Tensor::Uniform(float from, float to, std::shared_ptr generator) { return nn::init::Uniform(shared_from_this(), from, to, generator); } diff --git a/scripts/check_generator_reproducibility.sh b/scripts/check_generator_reproducibility.sh new file mode 100755 index 00000000..b447e657 --- /dev/null +++ b/scripts/check_generator_reproducibility.sh @@ -0,0 +1,29 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd) +test_binary=${1:-"$repo_root/build-cpu/tests/tensor/test_tensor_cpu_only"} + +if [[ ! -x "$test_binary" ]]; then + echo "Generator test binary is not executable: $test_binary" >&2 + exit 1 +fi + +run_digest() { + "$test_binary" --gtest_filter='CPUGeneratorTest.CrossProcessReproducibilityDigest' 2>&1 \ + | sed -n 's/^GENERATOR_REPRODUCIBILITY_DIGEST=//p' +} + +first=$(run_digest) +second=$(run_digest) + +if [[ -z "$first" || -z "$second" ]]; then + echo "Failed to capture reproducibility digest" >&2 + exit 1 +fi +if [[ "$first" != "$second" ]]; then + echo "Cross-process reproducibility mismatch: $first != $second" >&2 + exit 1 +fi + +echo "Cross-process Generator reproducibility passed: $first" diff --git a/tests/tensor/cpu_only/test_generator.cc b/tests/tensor/cpu_only/test_generator.cc new file mode 100644 index 00000000..42b291c7 --- /dev/null +++ b/tests/tensor/cpu_only/test_generator.cc @@ -0,0 +1,663 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/generator.h" +#include "infini_train/include/nn/functional.h" +#include "infini_train/include/nn/init.h" +#include "infini_train/include/tensor.h" + +using namespace infini_train; + +namespace { +std::vector TensorData(const std::shared_ptr &tensor) { + CHECK(tensor->GetDevice().IsCPU()); + if (tensor->NumElements() == 0) { + return {}; + } + const auto *data = static_cast(tensor->DataPtr()); + return std::vector(data, data + tensor->NumElements()); +} + +bool SameValues(const std::vector &a, const std::vector &b) { return a == b; } + +bool DifferentValues(const std::vector &a, const std::vector &b) { return a != b; } + +uint64_t HashValues(uint64_t hash, const std::vector &values) { + constexpr uint64_t kFnvPrime = 1099511628211ULL; + for (float value : values) { + const uint32_t bits = std::bit_cast(value); + for (int shift = 0; shift < 32; shift += 8) { + hash ^= (bits >> shift) & 0xffU; + hash *= kFnvPrime; + } + } + return hash; +} + +std::string SerializedState(const std::vector &state) { return std::string(state.begin(), state.end()); } + +std::vector StateBytes(const std::string &state) { return std::vector(state.begin(), state.end()); } + +uint64_t CUDASemanticOffset(const std::vector &state) { + std::istringstream iss(SerializedState(state)); + std::string line; + CHECK(std::getline(iss, line)); + CHECK(std::getline(iss, line)); + CHECK(std::getline(iss, line)); + CHECK(std::getline(iss, line)); + return std::stoull(line); +} + +std::vector ReplaceStateLine(const std::vector &state, size_t line_index, + const std::string &replacement) { + std::string serialized = SerializedState(state); + size_t offset_begin = 0; + for (size_t i = 0; i < line_index; ++i) { + offset_begin = serialized.find('\n', offset_begin); + CHECK_NE(offset_begin, std::string::npos); + ++offset_begin; + } + size_t offset_end = serialized.find('\n', offset_begin); + if (offset_end == std::string::npos) { + offset_end = serialized.size(); + } + serialized.replace(offset_begin, offset_end - offset_begin, replacement); + return StateBytes(serialized); +} + +std::vector ReplaceCPUEngineToken(const std::vector &state, size_t token_index, + const std::string &replacement) { + const std::string serialized = SerializedState(state); + size_t engine_begin = serialized.find('\n'); + CHECK_NE(engine_begin, std::string::npos); + engine_begin = serialized.find('\n', engine_begin + 1); + CHECK_NE(engine_begin, std::string::npos); + ++engine_begin; + + std::istringstream engine_stream(serialized.substr(engine_begin)); + std::vector tokens; + for (std::string token; engine_stream >> token;) { tokens.push_back(token); } + CHECK_LT(token_index, tokens.size()); + tokens[token_index] = replacement; + + std::ostringstream rebuilt; + rebuilt << serialized.substr(0, engine_begin); + for (size_t i = 0; i < tokens.size(); ++i) { + if (i != 0) { + rebuilt << ' '; + } + rebuilt << tokens[i]; + } + return StateBytes(rebuilt.str()); +} + +std::vector MakeAllZeroCPUEngineState(const std::vector &state) { + const std::string serialized = SerializedState(state); + size_t engine_begin = serialized.find('\n'); + CHECK_NE(engine_begin, std::string::npos); + engine_begin = serialized.find('\n', engine_begin + 1); + CHECK_NE(engine_begin, std::string::npos); + ++engine_begin; + + std::ostringstream rebuilt; + rebuilt << serialized.substr(0, engine_begin); + for (size_t i = 0; i < std::mt19937::state_size; ++i) { + if (i != 0) { + rebuilt << ' '; + } + rebuilt << '0'; + } + rebuilt << ' ' << std::mt19937::state_size; + return StateBytes(rebuilt.str()); +} + +std::vector ReplaceCUDASemanticOffset(const std::vector &state, const std::string &offset) { + return ReplaceStateLine(state, 3, offset); +} +} // namespace + +TEST(CPUGeneratorTest, InterfaceAndStateRoundTrip) { + auto generator = MakeCPUGenerator(123); + + EXPECT_EQ(generator->InitialSeed(), 123U); + EXPECT_TRUE(generator->GetDevice().IsCPU()); + + auto state = generator->GetState(); + EXPECT_FALSE(state.empty()); + + generator->ManualSeed(456); + EXPECT_EQ(generator->InitialSeed(), 456U); + + const auto generated_seed = generator->Seed(); + EXPECT_EQ(generator->InitialSeed(), generated_seed); + + generator->SetState(state); + EXPECT_EQ(generator->InitialSeed(), 123U); +} + +TEST(CPUGeneratorTest, ProcessDefaultSeedIsSharedByInitialDefaultGenerators) { + auto cpu = GetDefaultCPUGenerator(); + auto cuda = GetDefaultCUDAGenerator(0); + + EXPECT_EQ(cpu->InitialSeed(), cuda->InitialSeed()); +} + +TEST(CPUGeneratorTest, RandUsesManualSeedForReproducibility) { + ManualSeed(123); + auto a = TensorData(nn::function::Rand({2, 3})); + + ManualSeed(123); + auto b = TensorData(nn::function::Rand({2, 3})); + + ManualSeed(456); + auto c = TensorData(nn::function::Rand({2, 3})); + + EXPECT_TRUE(SameValues(a, b)); + EXPECT_TRUE(DifferentValues(a, c)); +} + +TEST(CPUGeneratorTest, RandnUsesManualSeedForReproducibility) { + ManualSeed(123); + auto a = TensorData(nn::function::Randn({2, 3})); + + ManualSeed(123); + auto b = TensorData(nn::function::Randn({2, 3})); + + ManualSeed(456); + auto c = TensorData(nn::function::Randn({2, 3})); + + EXPECT_TRUE(SameValues(a, b)); + EXPECT_TRUE(DifferentValues(a, c)); +} + +TEST(CPUGeneratorTest, CrossProcessReproducibilityDigest) { + ManualSeed(20260712); + auto parameter = std::make_shared(std::vector{8, 16}, DataType::kFLOAT32, Device()); + auto initialized = TensorData(nn::init::KaimingUniform(parameter)); + auto uniform = TensorData(nn::function::Rand({31})); + auto normal = TensorData(nn::function::Randn({29})); + + uint64_t digest = 1469598103934665603ULL; + digest = HashValues(digest, initialized); + digest = HashValues(digest, uniform); + digest = HashValues(digest, normal); + std::cout << "GENERATOR_REPRODUCIBILITY_DIGEST=" << std::hex << digest << std::dec << std::endl; +} + +TEST(CPUGeneratorTest, ModelLevelInitializationAndRandomSequenceReplay) { + ManualSeed(2026); + auto parameter = std::make_shared(std::vector{4, 8}, DataType::kFLOAT32, Device()); + auto initialized = TensorData(nn::init::KaimingUniform(parameter)); + auto state = GetDefaultCPUGenerator()->GetState(); + auto uniform = TensorData(nn::function::Rand({17})); + auto normal = TensorData(nn::function::Randn({19})); + + GetDefaultCPUGenerator()->SetState(state); + EXPECT_EQ(uniform, TensorData(nn::function::Rand({17}))); + EXPECT_EQ(normal, TensorData(nn::function::Randn({19}))); + + ManualSeed(2026); + auto replay_parameter = std::make_shared(std::vector{4, 8}, DataType::kFLOAT32, Device()); + EXPECT_EQ(initialized, TensorData(nn::init::KaimingUniform(replay_parameter))); +} + +TEST(CPUGeneratorTest, ManualSeedUsesHighBitsOfUint64Seed) { + constexpr uint64_t low_seed = 123; + constexpr uint64_t high_seed = low_seed + (1ULL << 32); + + auto cpu_low = MakeCPUGenerator(low_seed); + auto cpu_high = MakeCPUGenerator(high_seed); + + auto cpu_a = TensorData(nn::function::Rand({8}, Device(), cpu_low)); + auto cpu_b = TensorData(nn::function::Rand({8}, Device(), cpu_high)); + + EXPECT_EQ(cpu_high->InitialSeed(), high_seed); + EXPECT_TRUE(DifferentValues(cpu_a, cpu_b)); + + auto cuda_low = MakeCUDAGenerator(0, low_seed); + auto cuda_high = MakeCUDAGenerator(0, high_seed); + std::vector cuda_a(8); + std::vector cuda_b(8); + cuda_low->FillUniform(cuda_a, 0.0f, 1.0f); + cuda_high->FillUniform(cuda_b, 0.0f, 1.0f); + + EXPECT_EQ(cuda_high->InitialSeed(), high_seed); + EXPECT_TRUE(DifferentValues(cuda_a, cuda_b)); +} + +TEST(CPUGeneratorTest, StateRestoreReplaysSequence) { + auto generator = MakeCPUGenerator(123); + + (void)TensorData(nn::function::Rand({4}, Device(), generator)); + auto state = generator->GetState(); + + auto b = TensorData(nn::function::Rand({4}, Device(), generator)); + generator->SetState(state); + auto b2 = TensorData(nn::function::Rand({4}, Device(), generator)); + + EXPECT_TRUE(SameValues(b, b2)); +} + +TEST(CPUGeneratorTest, SetStateRejectsCorruptedState) { + auto generator = MakeCPUGenerator(123); + std::vector corrupted_state = {'n', 'o', 't', '-', 'a', '-', 's', 't', 'a', 't', 'e'}; + + EXPECT_DEATH(generator->SetState(corrupted_state), "Invalid CPU generator state header"); +} + +TEST(CPUGeneratorTest, SetStateRejectsEmptyAndTruncatedState) { + auto cpu_generator = MakeCPUGenerator(123); + EXPECT_DEATH(cpu_generator->SetState({}), "Invalid CPU generator state header"); + + auto cpu_state = cpu_generator->GetState(); + cpu_state.resize(cpu_state.size() / 2); + EXPECT_DEATH(cpu_generator->SetState(cpu_state), "Invalid CPU generator engine state"); + + auto cuda_generator = MakeCUDAGenerator(0, 123); + auto cuda_state = cuda_generator->GetState(); + const auto offset_separator = SerializedState(cuda_state).rfind('\n', cuda_state.size() - 2); + ASSERT_NE(offset_separator, std::string::npos); + cuda_state.resize(offset_separator + 1); + EXPECT_DEATH(cuda_generator->SetState(cuda_state), "Invalid CUDA generator offset in state"); +} + +TEST(CPUGeneratorTest, SetStateRejectsTrailingStateData) { + auto cpu_generator = MakeCPUGenerator(123); + auto cpu_state = cpu_generator->GetState(); + const std::string trailing_data = " trailing-junk"; + cpu_state.insert(cpu_state.end(), trailing_data.begin(), trailing_data.end()); + + EXPECT_DEATH(cpu_generator->SetState(cpu_state), "trailing data"); + + auto cuda_generator = MakeCUDAGenerator(0, 123); + auto cuda_state = cuda_generator->GetState(); + cuda_state.insert(cuda_state.end(), trailing_data.begin(), trailing_data.end()); + + EXPECT_DEATH(cuda_generator->SetState(cuda_state), "trailing data"); +} + +TEST(CPUGeneratorTest, SetStateRejectsTrailingWhitespace) { + auto cpu_generator = MakeCPUGenerator(123); + auto cpu_state = cpu_generator->GetState(); + cpu_state.push_back(' '); + EXPECT_DEATH(cpu_generator->SetState(cpu_state), "trailing data"); + + auto cuda_generator = MakeCUDAGenerator(0, 123); + auto cuda_state = cuda_generator->GetState(); + cuda_state.push_back('\n'); + EXPECT_DEATH(cuda_generator->SetState(cuda_state), "trailing data"); +} + +TEST(CPUGeneratorTest, SetStateRejectsWrongBackendState) { + auto cpu_generator = MakeCPUGenerator(123); + auto cuda_generator = MakeCUDAGenerator(0, 123); + + EXPECT_DEATH(cpu_generator->SetState(cuda_generator->GetState()), "Invalid CPU generator state header"); + EXPECT_DEATH(cuda_generator->SetState(cpu_generator->GetState()), "Invalid CUDA generator state header"); +} + +TEST(CPUGeneratorTest, SetStateRejectsNegativeUnsignedFields) { + auto cpu_generator = MakeCPUGenerator(123); + auto negative_cpu_seed = ReplaceStateLine(cpu_generator->GetState(), 1, "-1"); + EXPECT_DEATH(cpu_generator->SetState(negative_cpu_seed), "Invalid CPU generator seed in state"); + + auto cuda_generator = MakeCUDAGenerator(0, 123); + auto negative_cuda_seed = ReplaceStateLine(cuda_generator->GetState(), 1, "-1"); + EXPECT_DEATH(cuda_generator->SetState(negative_cuda_seed), "Invalid CUDA generator seed in state"); + + auto negative_cuda_offset = ReplaceStateLine(cuda_generator->GetState(), 3, "-1"); + EXPECT_DEATH(cuda_generator->SetState(negative_cuda_offset), "Invalid CUDA generator offset in state"); +} + +TEST(CPUGeneratorTest, CPUSetStateRejectsInvalidEngineStructure) { + auto generator = MakeCPUGenerator(123); + const auto state = generator->GetState(); + + auto oversized_word = ReplaceCPUEngineToken(state, 0, "4294967296"); + EXPECT_DEATH(generator->SetState(oversized_word), "Invalid CPU generator engine word in state"); + + auto invalid_position + = ReplaceCPUEngineToken(state, std::mt19937::state_size, std::to_string(std::mt19937::state_size + 1)); + EXPECT_DEATH(generator->SetState(invalid_position), "Invalid CPU generator engine position in state"); + + auto all_zero = MakeAllZeroCPUEngineState(state); + EXPECT_DEATH(generator->SetState(all_zero), "Invalid CPU generator engine state: all-zero state"); +} + +TEST(CPUGeneratorTest, CUDAGeneratorSetStateRejectsInvalidSourceDeviceIndex) { + auto generator = MakeCUDAGenerator(0, 123); + auto negative_index = ReplaceStateLine(generator->GetState(), 2, "-1"); + EXPECT_DEATH(generator->SetState(negative_index), "Invalid CUDA generator device index in state"); + + auto oversized_index = ReplaceStateLine(generator->GetState(), 2, "128"); + EXPECT_DEATH(generator->SetState(oversized_index), "Invalid CUDA generator device index in state"); +} + +TEST(CPUGeneratorTest, CUDAGeneratorSetStateFromDifferentDeviceReplaysSequence) { + auto cuda0 = MakeCUDAGenerator(0, 123); + std::vector ignored(5); + std::vector expected(11); + std::vector actual(11); + cuda0->FillUniform(ignored, 0.0f, 1.0f); + auto state = cuda0->GetState(); + cuda0->FillUniform(expected, 0.0f, 1.0f); + + auto cuda1 = MakeCUDAGenerator(1, 999); + cuda1->SetState(state); + cuda1->FillUniform(actual, 0.0f, 1.0f); + + EXPECT_EQ(expected, actual); + EXPECT_EQ(cuda1->InitialSeed(), 123U); + EXPECT_EQ(cuda1->GetDevice().index(), 1); +} + +TEST(CPUGeneratorTest, DefaultGeneratorAdvancesAcrossCalls) { + ManualSeed(123); + auto a = TensorData(nn::function::Rand({4})); + auto b = TensorData(nn::function::Rand({4})); + + ManualSeed(123); + auto a2 = TensorData(nn::function::Rand({4})); + + EXPECT_TRUE(DifferentValues(a, b)); + EXPECT_TRUE(SameValues(a, a2)); +} + +TEST(CPUGeneratorTest, ExplicitGeneratorDoesNotAdvanceDefaultGenerator) { + ManualSeed(777); + auto default_generator = GetDefaultCPUGenerator(); + auto before = default_generator->GetState(); + + auto explicit_generator = MakeCPUGenerator(123); + (void)TensorData(nn::function::Rand({8}, Device(), explicit_generator)); + + auto after = default_generator->GetState(); + EXPECT_EQ(before, after); +} + +TEST(CPUGeneratorTest, TensorUniformAcceptsExplicitGenerator) { + auto g1 = MakeCPUGenerator(999); + auto t1 = std::make_shared(std::vector{2, 3}, DataType::kFLOAT32, Device()); + auto a = TensorData(t1->Uniform(-1.0f, 1.0f, g1)); + + auto g2 = MakeCPUGenerator(999); + auto t2 = std::make_shared(std::vector{2, 3}, DataType::kFLOAT32, Device()); + auto b = TensorData(t2->Uniform(-1.0f, 1.0f, g2)); + + EXPECT_TRUE(SameValues(a, b)); + EXPECT_TRUE(std::all_of(a.begin(), a.end(), [](float x) { return x >= -1.0f && x <= 1.0f; })); +} + +TEST(CPUGeneratorTest, ZeroElementRandomCallsDoNotAdvanceState) { + auto generator = MakeCPUGenerator(123); + auto before = generator->GetState(); + EXPECT_EQ(TensorData(nn::function::Rand({0}, Device(), generator)).size(), 0U); + EXPECT_EQ(TensorData(nn::function::Randn({0}, Device(), generator)).size(), 0U); + EXPECT_EQ(before, generator->GetState()); +} + +TEST(CPUGeneratorTest, RandomDistributionsRejectInvalidParameters) { + auto tensor = std::make_shared(std::vector{4}, DataType::kFLOAT32, Device()); + EXPECT_DEATH((void)nn::init::Uniform(tensor, 2.0f, 1.0f), "Check failed: from <= to"); + EXPECT_DEATH((void)nn::init::Normal(tensor, 0.0f, -0.1f), "Check failed: std >= 0.0f"); +} + +TEST(CPUGeneratorTest, EmptyRandomCallsStillValidateParametersAndBackend) { + auto tensor = std::make_shared(std::vector{0}, DataType::kFLOAT32, Device()); + EXPECT_DEATH((void)nn::init::Normal(tensor, 0.0f, -0.1f), "Check failed: std >= 0.0f"); + + auto cuda_generator = MakeCUDAGenerator(0, 123); + EXPECT_DEATH((void)nn::init::Uniform(tensor, 0.0f, 1.0f, cuda_generator), + "Generator backend must match tensor device backend"); +} + +TEST(CPUGeneratorTest, UniformRejectsNonFiniteAndOverflowingBoundsEvenWhenEmpty) { + auto tensor = std::make_shared(std::vector{0}, DataType::kFLOAT32, Device()); + const float infinity = std::numeric_limits::infinity(); + const float max = std::numeric_limits::max(); + EXPECT_DEATH((void)nn::init::Uniform(tensor, 0.0f, infinity), "Uniform upper bound must be finite"); + EXPECT_DEATH((void)nn::init::Uniform(tensor, -max, max), "Uniform bounds range exceeds float maximum"); + + std::vector empty; + auto cpu_generator = MakeCPUGenerator(123); + EXPECT_DEATH(cpu_generator->FillUniform(empty, 0.0f, infinity), "Uniform upper bound must be finite"); + + auto cuda_generator = MakeCUDAGenerator(0, 123); + EXPECT_DEATH(cuda_generator->FillUniform(empty, -max, max), "Uniform bounds range exceeds float maximum"); +} + +TEST(CPUGeneratorTest, EqualUniformBoundsReturnConstantAndAdvanceState) { + constexpr size_t kNumElements = 17; + auto generator = MakeCPUGenerator(123); + auto tensor = std::make_shared(std::vector{kNumElements}, DataType::kFLOAT32, Device()); + const auto before = generator->GetState(); + + auto values = TensorData(nn::init::Uniform(tensor, 2.5f, 2.5f, generator)); + + EXPECT_TRUE(std::all_of(values.begin(), values.end(), [](float value) { return value == 2.5f; })); + EXPECT_NE(before, generator->GetState()); +} + +TEST(CPUGeneratorTest, RandomInitializationRejectsNonFloat32Tensors) { + auto tensor = std::make_shared(std::vector{4}, DataType::kFLOAT16, Device()); + EXPECT_DEATH((void)nn::init::Uniform(tensor), "Random uniform currently supports float32 tensors"); + EXPECT_DEATH((void)nn::init::Normal(tensor), "Random normal currently supports float32 tensors"); +} + +TEST(CPUGeneratorTest, ZeroStdNormalReturnsMeanAndAdvancesState) { + constexpr size_t kNumElements = 17; + auto zero_std_generator = MakeCPUGenerator(123); + auto reference_generator = MakeCPUGenerator(123); + auto tensor = std::make_shared(std::vector{kNumElements}, DataType::kFLOAT32, Device()); + + auto values = TensorData(nn::init::Normal(tensor, 2.5f, 0.0f, zero_std_generator)); + std::vector ignored(kNumElements); + reference_generator->FillNormal(ignored, 0.0f, 1.0f); + + EXPECT_TRUE(std::all_of(values.begin(), values.end(), [](float value) { return value == 2.5f; })); + EXPECT_EQ(zero_std_generator->GetState(), reference_generator->GetState()); +} + +TEST(CPUGeneratorTest, CUDAHostZeroStdNormalReturnsMeanAndAdvancesOffset) { + constexpr size_t kNumElements = 17; + auto generator = MakeCUDAGenerator(0, 123); + std::vector values(kNumElements); + + generator->FillNormal(values, 2.5f, 0.0f); + + EXPECT_TRUE(std::all_of(values.begin(), values.end(), [](float value) { return value == 2.5f; })); + EXPECT_EQ(CUDASemanticOffset(generator->GetState()), kNumElements * 2); +} + +TEST(CPUGeneratorTest, LargeRandomTensorSpansManyBlocksWorthOfValues) { + auto values = TensorData(nn::function::Rand({65537}, Device(), MakeCPUGenerator(123))); + EXPECT_EQ(values.size(), 65537U); + EXPECT_TRUE(std::all_of(values.begin(), values.end(), [](float value) { return value >= 0.0f && value <= 1.0f; })); +} + +TEST(CPUGeneratorTest, KaimingUniformUsesExplicitGeneratorReproducibly) { + auto g1 = MakeCPUGenerator(314159); + auto t1 = std::make_shared(std::vector{4, 8}, DataType::kFLOAT32, Device()); + auto a = TensorData( + nn::init::KaimingUniform(t1, 0.0f, nn::init::KaimingMode::kFanIn, nn::init::NonLinearityType::kReLU, g1)); + + auto g2 = MakeCPUGenerator(314159); + auto t2 = std::make_shared(std::vector{4, 8}, DataType::kFLOAT32, Device()); + auto b = TensorData( + nn::init::KaimingUniform(t2, 0.0f, nn::init::KaimingMode::kFanIn, nn::init::NonLinearityType::kReLU, g2)); + + EXPECT_EQ(a, b); +} + +TEST(CPUGeneratorTest, ExplicitGeneratorBackendMustMatchTensorDeviceBackend) { + auto tensor = std::make_shared(std::vector{2, 3}, DataType::kFLOAT32, Device()); + auto cuda_generator = MakeCUDAGenerator(0, 123); + + EXPECT_DEATH((void)tensor->Uniform(0.0f, 1.0f, cuda_generator), + "Generator backend must match tensor device backend"); +} + +TEST(CPUGeneratorTest, DefaultCUDAGeneratorsArePerDevice) { + ManualSeedAll(321); + auto cuda0 = GetDefaultCUDAGenerator(0); + auto cuda1 = GetDefaultCUDAGenerator(1); + + EXPECT_TRUE(cuda0->GetDevice().IsCUDA()); + EXPECT_TRUE(cuda1->GetDevice().IsCUDA()); + EXPECT_EQ(cuda0->GetDevice().index(), 0); + EXPECT_EQ(cuda1->GetDevice().index(), 1); + EXPECT_EQ(cuda0->InitialSeed(), 321U); + EXPECT_EQ(cuda1->InitialSeed(), 321U); + EXPECT_NE(cuda0->GetState(), cuda1->GetState()); + + std::vector buffer(4); + auto cuda1_before = cuda1->GetState(); + cuda0->FillUniform(buffer, 0.0f, 1.0f); + EXPECT_EQ(cuda1_before, cuda1->GetState()); + EXPECT_EQ(cuda0, GetDefaultCUDAGenerator(0)); +} + +TEST(CPUGeneratorTest, GetDefaultGeneratorDispatchesByDevice) { + ManualSeedAll(654); + + auto cpu = GetDefaultGenerator(Device()); + auto cuda0 = GetDefaultGenerator(Device(Device::DeviceType::kCUDA, 0)); + + EXPECT_EQ(cpu, GetDefaultCPUGenerator()); + EXPECT_EQ(cuda0, GetDefaultCUDAGenerator(0)); + EXPECT_TRUE(cpu->GetDevice().IsCPU()); + EXPECT_TRUE(cuda0->GetDevice().IsCUDA()); + EXPECT_EQ(cuda0->GetDevice().index(), 0); + EXPECT_EQ(cpu->InitialSeed(), 654U); + EXPECT_EQ(cuda0->InitialSeed(), 654U); +} + +TEST(CPUGeneratorTest, ManualSeedAllSeedsExistingAndFutureDefaultGenerators) { + auto cuda0 = GetDefaultCUDAGenerator(0); + cuda0->ManualSeed(123); + + ManualSeedAll(888); + + EXPECT_EQ(GetDefaultCPUGenerator()->InitialSeed(), 888U); + EXPECT_EQ(cuda0->InitialSeed(), 888U); + EXPECT_EQ(GetDefaultCUDAGenerator(2)->InitialSeed(), 888U); +} + +TEST(CPUGeneratorTest, ManualSeedAliasesManualSeedAll) { + auto cuda0 = GetDefaultCUDAGenerator(0); + cuda0->ManualSeed(123); + + ManualSeed(999); + + EXPECT_EQ(GetDefaultCPUGenerator()->InitialSeed(), 999U); + EXPECT_EQ(cuda0->InitialSeed(), 999U); + EXPECT_EQ(GetDefaultCUDAGenerator(3)->InitialSeed(), 999U); +} + +TEST(CPUGeneratorTest, CUDAGeneratorStateRestoreReplaysSequence) { + auto generator = MakeCUDAGenerator(0, 123); + std::vector ignored(4); + std::vector b(4); + std::vector b2(4); + + generator->FillUniform(ignored, 0.0f, 1.0f); + auto state = generator->GetState(); + generator->FillUniform(b, 0.0f, 1.0f); + generator->SetState(state); + generator->FillUniform(b2, 0.0f, 1.0f); + + EXPECT_TRUE(SameValues(b, b2)); +} + +TEST(CPUGeneratorTest, CUDAGeneratorStateSerializesSemanticOffset) { + auto generator = MakeCUDAGenerator(0, 123); + EXPECT_EQ(CUDASemanticOffset(generator->GetState()), 0U); + + std::vector uniform(5); + generator->FillUniform(uniform, 0.0f, 1.0f); + EXPECT_EQ(CUDASemanticOffset(generator->GetState()), 5U); + + std::vector normal(7); + generator->FillNormal(normal, 0.0f, 1.0f); + auto state = generator->GetState(); + EXPECT_EQ(CUDASemanticOffset(state), 19U); + + generator->ManualSeed(456); + EXPECT_EQ(CUDASemanticOffset(generator->GetState()), 0U); + + generator->SetState(state); + EXPECT_EQ(CUDASemanticOffset(generator->GetState()), 19U); +} + +TEST(CPUGeneratorTest, CUDAGeneratorSetStateRejectsCorruptedOffset) { + auto generator = MakeCUDAGenerator(0, 123); + auto state = ReplaceCUDASemanticOffset(generator->GetState(), "not-an-offset"); + + EXPECT_DEATH(generator->SetState(state), "Invalid CUDA generator offset in state"); +} + +TEST(CPUGeneratorTest, CUDAGeneratorRejectsOffsetOverflow) { + auto generator = MakeCUDAGenerator(0, 123); + EXPECT_EQ(generator->ReserveRandomOffset(std::numeric_limits::max()).second, 0U); + EXPECT_DEATH((void)generator->ReserveRandomOffset(1), "CUDA generator offset overflow"); +} + +TEST(CPUGeneratorTest, CUDAGeneratorOffsetReservationsAreThreadSafeAndDisjoint) { + constexpr size_t kThreadCount = 8; + constexpr size_t kReservationsPerThread = 100; + constexpr uint64_t kIncrement = 7; + auto generator = MakeCUDAGenerator(0, 123); + std::vector offsets(kThreadCount * kReservationsPerThread); + std::vector threads; + threads.reserve(kThreadCount); + + for (size_t thread_index = 0; thread_index < kThreadCount; ++thread_index) { + threads.emplace_back([&, thread_index]() { + for (size_t i = 0; i < kReservationsPerThread; ++i) { + const auto [seed, offset] = generator->ReserveRandomOffset(kIncrement); + EXPECT_EQ(seed, 123U); + offsets[thread_index * kReservationsPerThread + i] = offset; + } + }); + } + for (auto &thread : threads) { thread.join(); } + + std::sort(offsets.begin(), offsets.end()); + for (size_t i = 0; i < offsets.size(); ++i) { EXPECT_EQ(offsets[i], static_cast(i) * kIncrement); } + EXPECT_EQ(CUDASemanticOffset(generator->GetState()), offsets.size() * kIncrement); +} + +TEST(CPUGeneratorTest, DefaultCUDAGeneratorIsThreadSafeUnderConcurrentConsumption) { + constexpr size_t kThreadCount = 8; + constexpr size_t kReservationsPerThread = 50; + constexpr uint64_t kIncrement = 3; + ManualSeedAll(4242); + auto generator = GetDefaultCUDAGenerator(0); + std::vector offsets(kThreadCount * kReservationsPerThread); + std::vector threads; + + for (size_t thread_index = 0; thread_index < kThreadCount; ++thread_index) { + threads.emplace_back([&, thread_index]() { + for (size_t i = 0; i < kReservationsPerThread; ++i) { + offsets[thread_index * kReservationsPerThread + i] = generator->ReserveRandomOffset(kIncrement).second; + } + }); + } + for (auto &thread : threads) { thread.join(); } + + std::sort(offsets.begin(), offsets.end()); + for (size_t i = 0; i < offsets.size(); ++i) { EXPECT_EQ(offsets[i], static_cast(i) * kIncrement); } + EXPECT_EQ(CUDASemanticOffset(generator->GetState()), offsets.size() * kIncrement); +} diff --git a/tests/tensor/cuda_only/test_generator_cuda.cc b/tests/tensor/cuda_only/test_generator_cuda.cc new file mode 100644 index 00000000..05881cde --- /dev/null +++ b/tests/tensor/cuda_only/test_generator_cuda.cc @@ -0,0 +1,314 @@ +#include +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/generator.h" +#include "infini_train/include/nn/functional.h" +#include "infini_train/include/nn/init.h" +#include "infini_train/include/tensor.h" + +#include "tests/common/test_utils.h" + +using namespace infini_train; + +namespace { +std::vector CopyToCPUData(const std::shared_ptr &tensor) { + auto cpu_tensor = tensor->To(Device()); + auto *impl = core::GetDeviceGuardImpl(tensor->GetDevice().type()); + impl->SynchronizeDevice(tensor->GetDevice()); + + if (cpu_tensor.NumElements() == 0) { + return {}; + } + const auto *data = static_cast(cpu_tensor.DataPtr()); + return std::vector(data, data + cpu_tensor.NumElements()); +} + +uint64_t CUDASemanticOffset(const std::vector &state) { + std::istringstream iss(std::string(state.begin(), state.end())); + std::string line; + CHECK(std::getline(iss, line)); + CHECK(std::getline(iss, line)); + CHECK(std::getline(iss, line)); + CHECK(std::getline(iss, line)); + return std::stoull(line); +} +} // namespace + +class CUDAGeneratorTensorTest : public ::testing::Test {}; + +TEST_F(CUDAGeneratorTensorTest, DefaultCUDAGeneratorControlsRand) { + REQUIRE_MIN_DEVICES(1); + Device cuda_device(Device::DeviceType::kCUDA, 0); + + ManualSeedAll(123); + auto a = CopyToCPUData(nn::function::Rand({2, 3}, cuda_device)); + + ManualSeedAll(123); + auto b = CopyToCPUData(nn::function::Rand({2, 3}, cuda_device)); + + ManualSeedAll(456); + auto c = CopyToCPUData(nn::function::Rand({2, 3}, cuda_device)); + + EXPECT_EQ(a, b); + EXPECT_NE(a, c); +} + +TEST_F(CUDAGeneratorTensorTest, DefaultCUDAGeneratorControlsRandn) { + REQUIRE_MIN_DEVICES(1); + Device cuda_device(Device::DeviceType::kCUDA, 0); + + ManualSeedAll(123); + auto a = CopyToCPUData(nn::function::Randn({2, 3}, cuda_device)); + + ManualSeedAll(123); + auto b = CopyToCPUData(nn::function::Randn({2, 3}, cuda_device)); + + ManualSeedAll(456); + auto c = CopyToCPUData(nn::function::Randn({2, 3}, cuda_device)); + + EXPECT_EQ(a, b); + EXPECT_NE(a, c); +} + +TEST_F(CUDAGeneratorTensorTest, KaimingInitializationUsesDefaultGeneratorReproducibly) { + REQUIRE_MIN_DEVICES(1); + Device cuda_device(Device::DeviceType::kCUDA, 0); + + ManualSeedAll(314159); + auto first = std::make_shared(std::vector{8, 16}, DataType::kFLOAT32, cuda_device); + auto a = CopyToCPUData(nn::init::KaimingUniform(first)); + + ManualSeedAll(314159); + auto second = std::make_shared(std::vector{8, 16}, DataType::kFLOAT32, cuda_device); + auto b = CopyToCPUData(nn::init::KaimingUniform(second)); + + EXPECT_EQ(a, b); +} + +TEST_F(CUDAGeneratorTensorTest, DefaultCUDAGeneratorAdvancesAcrossRandCalls) { + REQUIRE_MIN_DEVICES(1); + Device cuda_device(Device::DeviceType::kCUDA, 0); + + ManualSeedAll(123); + auto a = CopyToCPUData(nn::function::Rand({4}, cuda_device)); + auto b = CopyToCPUData(nn::function::Rand({4}, cuda_device)); + + EXPECT_NE(a, b); +} + +TEST_F(CUDAGeneratorTensorTest, ExplicitCUDAGeneratorControlsRand) { + REQUIRE_MIN_DEVICES(1); + Device cuda_device(Device::DeviceType::kCUDA, 0); + + auto g1 = MakeCUDAGenerator(0, 123); + auto a = CopyToCPUData(nn::function::Rand({2, 3}, cuda_device, g1)); + + auto g2 = MakeCUDAGenerator(0, 123); + auto b = CopyToCPUData(nn::function::Rand({2, 3}, cuda_device, g2)); + + auto g3 = MakeCUDAGenerator(0, 456); + auto c = CopyToCPUData(nn::function::Rand({2, 3}, cuda_device, g3)); + + EXPECT_EQ(a, b); + EXPECT_NE(a, c); +} + +TEST_F(CUDAGeneratorTensorTest, CUDAGeneratorStateRestoreReplaysRand) { + REQUIRE_MIN_DEVICES(1); + Device cuda_device(Device::DeviceType::kCUDA, 0); + + auto generator = MakeCUDAGenerator(0, 123); + (void)CopyToCPUData(nn::function::Rand({4}, cuda_device, generator)); + auto state = generator->GetState(); + + auto b = CopyToCPUData(nn::function::Rand({4}, cuda_device, generator)); + generator->SetState(state); + auto b2 = CopyToCPUData(nn::function::Rand({4}, cuda_device, generator)); + + EXPECT_EQ(b, b2); +} + +TEST_F(CUDAGeneratorTensorTest, ExplicitCUDAGeneratorDoesNotAdvanceDefaultGenerator) { + REQUIRE_MIN_DEVICES(1); + Device cuda_device(Device::DeviceType::kCUDA, 0); + + ManualSeedAll(777); + auto default_generator = GetDefaultCUDAGenerator(0); + auto before = default_generator->GetState(); + + auto explicit_generator = MakeCUDAGenerator(0, 123); + (void)CopyToCPUData(nn::function::Rand({8}, cuda_device, explicit_generator)); + + auto after = default_generator->GetState(); + EXPECT_EQ(before, after); +} + +TEST_F(CUDAGeneratorTensorTest, TensorRandAdvancesCUDASemanticOffset) { + REQUIRE_MIN_DEVICES(1); + Device cuda_device(Device::DeviceType::kCUDA, 0); + + auto generator = MakeCUDAGenerator(0, 123); + EXPECT_EQ(CUDASemanticOffset(generator->GetState()), 0U); + + (void)CopyToCPUData(nn::function::Rand({2, 3}, cuda_device, generator)); + EXPECT_EQ(CUDASemanticOffset(generator->GetState()), 6U); + + (void)CopyToCPUData(nn::function::Randn({4}, cuda_device, generator)); + EXPECT_EQ(CUDASemanticOffset(generator->GetState()), 14U); +} + +TEST_F(CUDAGeneratorTensorTest, SplitCallsConsumeOneContinuousRandomStream) { + REQUIRE_MIN_DEVICES(1); + Device cuda_device(Device::DeviceType::kCUDA, 0); + + auto split_generator = MakeCUDAGenerator(0, 123); + auto first = CopyToCPUData(nn::function::Rand({13}, cuda_device, split_generator)); + auto second = CopyToCPUData(nn::function::Rand({19}, cuda_device, split_generator)); + first.insert(first.end(), second.begin(), second.end()); + + auto combined_generator = MakeCUDAGenerator(0, 123); + auto combined = CopyToCPUData(nn::function::Rand({32}, cuda_device, combined_generator)); + EXPECT_EQ(first, combined); + EXPECT_EQ(CUDASemanticOffset(split_generator->GetState()), 32U); + EXPECT_EQ(CUDASemanticOffset(combined_generator->GetState()), 32U); + + auto split_normal_generator = MakeCUDAGenerator(0, 456); + auto normal_first = CopyToCPUData(nn::function::Randn({7}, cuda_device, split_normal_generator)); + auto normal_second = CopyToCPUData(nn::function::Randn({11}, cuda_device, split_normal_generator)); + normal_first.insert(normal_first.end(), normal_second.begin(), normal_second.end()); + + auto combined_normal_generator = MakeCUDAGenerator(0, 456); + auto combined_normal = CopyToCPUData(nn::function::Randn({18}, cuda_device, combined_normal_generator)); + EXPECT_EQ(normal_first, combined_normal); + EXPECT_EQ(CUDASemanticOffset(split_normal_generator->GetState()), 36U); + EXPECT_EQ(CUDASemanticOffset(combined_normal_generator->GetState()), 36U); +} + +TEST_F(CUDAGeneratorTensorTest, HostAndKernelConsumptionShareOneOffsetStream) { + REQUIRE_MIN_DEVICES(1); + Device cuda_device(Device::DeviceType::kCUDA, 0); + + auto mixed_uniform_generator = MakeCUDAGenerator(0, 123); + std::vector host_uniform(13); + mixed_uniform_generator->FillUniform(host_uniform, 0.0f, 1.0f); + auto after_host_uniform = CopyToCPUData(nn::function::Rand({19}, cuda_device, mixed_uniform_generator)); + + auto reference_uniform_generator = MakeCUDAGenerator(0, 123); + (void)reference_uniform_generator->ReserveRandomOffset(13); + auto reference_uniform = CopyToCPUData(nn::function::Rand({19}, cuda_device, reference_uniform_generator)); + EXPECT_EQ(after_host_uniform, reference_uniform); + EXPECT_EQ(CUDASemanticOffset(mixed_uniform_generator->GetState()), 32U); + + auto mixed_normal_generator = MakeCUDAGenerator(0, 456); + std::vector host_normal(7); + mixed_normal_generator->FillNormal(host_normal, 0.0f, 1.0f); + auto after_host_normal = CopyToCPUData(nn::function::Randn({11}, cuda_device, mixed_normal_generator)); + + auto reference_normal_generator = MakeCUDAGenerator(0, 456); + (void)reference_normal_generator->ReserveRandomOffset(14); + auto reference_normal = CopyToCPUData(nn::function::Randn({11}, cuda_device, reference_normal_generator)); + EXPECT_EQ(after_host_normal, reference_normal); + EXPECT_EQ(CUDASemanticOffset(mixed_normal_generator->GetState()), 36U); +} + +TEST_F(CUDAGeneratorTensorTest, ZeroElementRandomCallsDoNotAdvanceState) { + REQUIRE_MIN_DEVICES(1); + Device cuda_device(Device::DeviceType::kCUDA, 0); + auto generator = MakeCUDAGenerator(0, 123); + auto before = generator->GetState(); + + EXPECT_TRUE(CopyToCPUData(nn::function::Rand({0}, cuda_device, generator)).empty()); + EXPECT_TRUE(CopyToCPUData(nn::function::Randn({0}, cuda_device, generator)).empty()); + EXPECT_EQ(before, generator->GetState()); +} + +TEST_F(CUDAGeneratorTensorTest, ZeroStdNormalReturnsMeanAndAdvancesState) { + REQUIRE_MIN_DEVICES(1); + Device cuda_device(Device::DeviceType::kCUDA, 0); + constexpr int64_t kNumElements = 17; + auto generator = MakeCUDAGenerator(0, 123); + auto tensor = std::make_shared(std::vector{kNumElements}, DataType::kFLOAT32, cuda_device); + + auto values = CopyToCPUData(nn::init::Normal(tensor, 2.5f, 0.0f, generator)); + + EXPECT_TRUE(std::all_of(values.begin(), values.end(), [](float value) { return value == 2.5f; })); + EXPECT_EQ(CUDASemanticOffset(generator->GetState()), static_cast(kNumElements) * 2); +} + +TEST_F(CUDAGeneratorTensorTest, EqualUniformBoundsReturnConstantAndAdvanceState) { + REQUIRE_MIN_DEVICES(1); + Device cuda_device(Device::DeviceType::kCUDA, 0); + constexpr int64_t kNumElements = 17; + auto generator = MakeCUDAGenerator(0, 123); + auto tensor = std::make_shared(std::vector{kNumElements}, DataType::kFLOAT32, cuda_device); + + auto values = CopyToCPUData(nn::init::Uniform(tensor, 2.5f, 2.5f, generator)); + + EXPECT_TRUE(std::all_of(values.begin(), values.end(), [](float value) { return value == 2.5f; })); + EXPECT_EQ(CUDASemanticOffset(generator->GetState()), static_cast(kNumElements)); +} + +TEST_F(CUDAGeneratorTensorTest, RandomKernelsHaveSaneRangeAndMoments) { + REQUIRE_MIN_DEVICES(1); + Device cuda_device(Device::DeviceType::kCUDA, 0); + constexpr int64_t kSampleCount = 65537; + + auto uniform = CopyToCPUData(nn::function::Rand({kSampleCount}, cuda_device, MakeCUDAGenerator(0, 2026))); + EXPECT_TRUE(std::all_of(uniform.begin(), uniform.end(), [](float value) { return value >= 0.0f && value < 1.0f; })); + const double uniform_mean = std::accumulate(uniform.begin(), uniform.end(), 0.0) / uniform.size(); + EXPECT_NEAR(uniform_mean, 0.5, 0.02); + + auto normal = CopyToCPUData(nn::function::Randn({kSampleCount}, cuda_device, MakeCUDAGenerator(0, 2026))); + EXPECT_TRUE(std::all_of(normal.begin(), normal.end(), [](float value) { return std::isfinite(value); })); + const double normal_mean = std::accumulate(normal.begin(), normal.end(), 0.0) / normal.size(); + const double squared_sum = std::inner_product(normal.begin(), normal.end(), normal.begin(), 0.0); + const double normal_variance = squared_sum / normal.size() - normal_mean * normal_mean; + EXPECT_NEAR(normal_mean, 0.0, 0.05); + EXPECT_NEAR(normal_variance, 1.0, 0.08); +} + +TEST_F(CUDAGeneratorTensorTest, ExplicitGeneratorMayTargetDifferentCUDAIndex) { + REQUIRE_MIN_DEVICES(2); + Device cuda1(Device::DeviceType::kCUDA, 1); + constexpr int64_t kNumElements = 32; + + auto generator = MakeCUDAGenerator(0, 123); + auto values = CopyToCPUData(nn::function::Rand({kNumElements}, cuda1, generator)); + + auto replay_generator = MakeCUDAGenerator(0, 123); + auto replay = CopyToCPUData(nn::function::Rand({kNumElements}, cuda1, replay_generator)); + + EXPECT_EQ(values, replay); + EXPECT_EQ(generator->GetDevice().index(), 0); + EXPECT_EQ(CUDASemanticOffset(generator->GetState()), static_cast(kNumElements)); +} + +TEST_F(CUDAGeneratorTensorTest, DefaultGeneratorsAreIndependentAcrossPhysicalDevices) { + REQUIRE_MIN_DEVICES(2); + Device cuda0(Device::DeviceType::kCUDA, 0); + Device cuda1(Device::DeviceType::kCUDA, 1); + + ManualSeedAll(2026); + auto generator0 = GetDefaultCUDAGenerator(0); + auto generator1 = GetDefaultCUDAGenerator(1); + + auto cuda1_before = generator1->GetState(); + auto a0 = CopyToCPUData(nn::function::Rand({32}, cuda0)); + EXPECT_EQ(cuda1_before, generator1->GetState()); + + auto cuda0_after = generator0->GetState(); + auto a1 = CopyToCPUData(nn::function::Rand({32}, cuda1)); + EXPECT_EQ(cuda0_after, generator0->GetState()); + EXPECT_EQ(a0, a1); + + auto b0 = CopyToCPUData(nn::function::Rand({32}, cuda0)); + EXPECT_NE(a0, b0); +}