diff --git a/infini_train/include/core/cpu_generator.h b/infini_train/include/core/cpu_generator.h new file mode 100644 index 00000000..d34d6329 --- /dev/null +++ b/infini_train/include/core/cpu_generator.h @@ -0,0 +1,57 @@ +#pragma once + +#include +#include +#include +#include + +#include "infini_train/include/core/generator.h" +#include "infini_train/include/device.h" + +namespace infini_train::core { + +/** + * CPU backend Generator. + * State organization: the random sequence is driven by a std::mt19937_64 + * engine. GetState() serializes a self-describing blob containing a backend + * magic, the current seed and the full engine state, allowing SetState() to + * resume the exact same sequence and to reject states coming from other + * backends. + */ +class CPUGeneratorImpl : public GeneratorImpl { +public: + static constexpr Device::DeviceType kDeviceType = Device::DeviceType::kCPU; + + explicit CPUGeneratorImpl(uint64_t seed = kDefaultSeed); + + void SetCurrentSeed(uint64_t seed) override; + uint64_t CurrentSeed() const override; + uint64_t Seed() override; + + // CPU mt19937 engine is not counter based; only offset 0 is supported. + void SetOffset(uint64_t offset) override; + uint64_t GetOffset() const override; + + std::vector GetState() const override; + void SetState(const std::vector &state) override; + + Device GetDevice() const override; + std::shared_ptr Clone() const override; + + // --- Internal kernel extension interface ---------------------------- + // The methods below expose backend engine details for use by random + // operators (e.g. init::Uniform). They are NOT part of the user-facing + // Generator API and should not be called from outside the kernel layer. + // + // Direct access to the underlying engine for random operators. Each draw + // advances the engine, i.e. advances the generator state. + std::mt19937_64 &Engine(); + +private: + static constexpr uint64_t kDefaultSeed = 67280421310721ULL; + + uint64_t seed_ = kDefaultSeed; + std::mt19937_64 engine_; +}; + +} // namespace infini_train::core diff --git a/infini_train/include/core/cuda_generator.h b/infini_train/include/core/cuda_generator.h new file mode 100644 index 00000000..d870aefc --- /dev/null +++ b/infini_train/include/core/cuda_generator.h @@ -0,0 +1,63 @@ +#pragma once + +#include +#include +#include +#include + +#include "infini_train/include/core/generator.h" +#include "infini_train/include/device.h" + +namespace infini_train::core { + +/** + * CUDA backend Generator. + * + * Unlike the CPU backend, CUDA random kernels are counter (Philox) based. The + * host side keeps only a lightweight state: the seed and a 64-bit Philox + * offset. Random kernels read the (seed, offset) pair and advance the offset by + * the number of random values they consumed, so that subsequent kernel launches + * do not overlap the random sequence. + * + * Each CUDA device maintains its own independent generator (see + * detail::DefaultCUDAGenerator). + */ +class CUDAGeneratorImpl : public GeneratorImpl { +public: + static constexpr Device::DeviceType kDeviceType = Device::DeviceType::kCUDA; + + explicit CUDAGeneratorImpl(int8_t device_index = 0, uint64_t seed = kDefaultSeed); + + void SetCurrentSeed(uint64_t seed) override; + uint64_t CurrentSeed() const override; + uint64_t Seed() override; + + void SetOffset(uint64_t offset) override; + uint64_t GetOffset() const override; + + std::vector GetState() const override; + void SetState(const std::vector &state) override; + + Device GetDevice() const override; + std::shared_ptr Clone() const override; + + // --- Internal kernel extension interface ---------------------------- + // The methods below expose Philox counter details for use by random + // operators (e.g. init::Uniform on CUDA tensors). They are NOT part of + // the user-facing Generator API and should not be called from outside the + // kernel layer. + // + // Reserves `increment` Philox values for the next kernel launch and returns + // the (seed, offset) pair the kernel must use. The offset is advanced so the + // next call yields a non-overlapping subsequence. + std::pair PhiloxEngineInputs(uint64_t increment); + +private: + static constexpr uint64_t kDefaultSeed = 67280421310721ULL; + + int8_t device_index_ = 0; + uint64_t seed_ = kDefaultSeed; + uint64_t philox_offset_ = 0; +}; + +} // namespace infini_train::core diff --git a/infini_train/include/core/generator.h b/infini_train/include/core/generator.h new file mode 100644 index 00000000..5c5bb46f --- /dev/null +++ b/infini_train/include/core/generator.h @@ -0,0 +1,171 @@ +#pragma once + +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/device.h" + +namespace infini_train::core { + +/** + * GeneratorImpl — Backend-agnostic polymorphic base class for random number generators. + */ + +class GeneratorImpl { +public: + GeneratorImpl() = default; + virtual ~GeneratorImpl() = default; + + // Delete copy and move in favor of Clone(). + GeneratorImpl(const GeneratorImpl &) = delete; + GeneratorImpl(GeneratorImpl &&) = delete; + GeneratorImpl &operator=(const GeneratorImpl &) = delete; + GeneratorImpl &operator=(GeneratorImpl &&) = delete; + + // Seed control -------------------------------------------------------- + virtual void SetCurrentSeed(uint64_t seed) = 0; + virtual uint64_t CurrentSeed() const = 0; + virtual uint64_t Seed() = 0; + // Philox-style offset control. Only meaningful for counter based backends + // (e.g. CUDA). Backends that cannot support random access may reject + // non-zero offsets. + virtual void SetOffset(uint64_t offset) = 0; + virtual uint64_t GetOffset() const = 0; + + // State control ------------------------------------------------------- + virtual std::vector GetState() const = 0; + virtual void SetState(const std::vector &state) = 0; + + virtual Device GetDevice() const = 0; + virtual std::shared_ptr Clone() const = 0; + + std::mutex mutex_; +}; + +/* + * Generator — User facing handle. Cheap to copy: copies share the same underlying impl, + * matching the semantics of a generator reference (e.g. torch.Generator). + */ +class Generator { +public: + Generator() = default; + + explicit Generator(std::shared_ptr impl) : impl_(std::move(impl)) { + CHECK(impl_ != nullptr) << "GeneratorImpl with nullptr is not supported"; + } + + bool Defined() const { return static_cast(impl_); } + + // Seed control -------------------------------------------------------- + void SetCurrentSeed(uint64_t seed) { + std::lock_guard lock(impl_->mutex_); + impl_->SetCurrentSeed(seed); + } + + void ManualSeed(uint64_t seed) { SetCurrentSeed(seed); } + uint64_t CurrentSeed() const { + std::lock_guard lock(impl_->mutex_); + return impl_->CurrentSeed(); + } + + uint64_t Seed() { + std::lock_guard lock(impl_->mutex_); + return impl_->Seed(); + } + + void SetOffset(uint64_t offset) { + std::lock_guard lock(impl_->mutex_); + impl_->SetOffset(offset); + } + uint64_t GetOffset() const { + std::lock_guard lock(impl_->mutex_); + return impl_->GetOffset(); + } + + // State control ------------------------------------------------------- + std::vector GetState() const { + std::lock_guard lock(impl_->mutex_); + return impl_->GetState(); + } + void SetState(const std::vector &state) { + std::lock_guard lock(impl_->mutex_); + impl_->SetState(state); + } + + // Misc ---------------------------------------------------------------- + Device GetDevice() const { return impl_->GetDevice(); } + std::mutex &Mutex() { return impl_->mutex_; } + + GeneratorImpl *UnsafeGetImpl() const { return impl_.get(); } + + template T *Get() const { return static_cast(impl_.get()); } + + Generator Clone() const { + std::lock_guard lock(impl_->mutex_); + return Generator(impl_->Clone()); + } + + bool operator==(const Generator &other) const { return impl_ == other.impl_; } + bool operator!=(const Generator &other) const { return !(*this == other); } + +private: + std::shared_ptr impl_; +}; + +namespace detail { +// Returns the process-wide default Generator for the given device. +const Generator &GetDefaultGenerator(const Device &device); + +// Convenience accessors. +const Generator &DefaultCPUGenerator(); +const Generator &DefaultCUDAGenerator(int8_t device_index = 0); + +// Appends a little-endian uint64_t to the output byte vector. +inline void AppendU64(std::vector &out, uint64_t value) { + for (int i = 0; i < 8; ++i) { out.push_back(static_cast((value >> (i * 8)) & 0xFF)); } +} + +// Reads a little-endian uint64_t from the given byte pointer. +inline uint64_t ReadU64(const uint8_t *p) { + uint64_t value = 0; + for (int i = 0; i < 8; ++i) { value |= static_cast(p[i]) << (i * 8); } + return value; +} +} // namespace detail + +// Global random seed entry point. +void ManualSeed(uint64_t seed); + +// Casts the given Generator (or the device default when none is supplied) to a +// concrete backend implementation, validating that the device type matches. +template T *GetGeneratorOrDefault(const std::optional &gen, const Device &device) { + const Generator &chosen = (gen.has_value() && gen->Defined()) ? *gen : detail::GetDefaultGenerator(device); + CHECK(chosen.Defined()) << "Generator with undefined implementation is not allowed"; + CHECK(T::kDeviceType == chosen.GetDevice().type()) + << "Generator device type mismatch: expected " << static_cast(T::kDeviceType) << " but found " + << static_cast(chosen.GetDevice().type()); + return chosen.Get(); +} + +// Convenience factory — matches PyTorch's at::make_generator(args...). +template Generator MakeGenerator(Args &&...args) { + return Generator(std::make_shared(std::forward(args)...)); +} + +// Validates a Generator optional and casts to the concrete backend type. +// Unlike GetGeneratorOrDefault, this does NOT fall back to a default — the +// caller must supply a defined generator. Matches PyTorch's check_generator. +template T *CheckGenerator(std::optional gen) { + CHECK(gen.has_value()) << "Expected a Generator but received std::nullopt"; + CHECK(gen->Defined()) << "Generator with undefined implementation is not allowed"; + CHECK(T::kDeviceType == gen->GetDevice().type()) + << "Generator device type mismatch: expected " << static_cast(T::kDeviceType) << " but found " + << static_cast(gen->GetDevice().type()); + return gen->Get(); +} + +} // namespace infini_train::core \ No newline at end of file diff --git a/infini_train/include/nn/init.h b/infini_train/include/nn/init.h index fc6effec..f820075f 100644 --- a/infini_train/include/nn/init.h +++ b/infini_train/include/nn/init.h @@ -2,9 +2,9 @@ #include #include -#include #include +#include "infini_train/include/core/generator.h" #include "infini_train/include/datatype.h" #include "infini_train/include/device.h" @@ -15,7 +15,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::optional generator = std::nullopt); std::pair CalculateFanInAndFanOut(const std::shared_ptr &tensor); @@ -42,10 +42,13 @@ 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::optional generator = std::nullopt); std::shared_ptr Uniform(const std::shared_ptr &tensor, float a = 0.0f, float b = 1.0f, - std::optional generator = std::nullopt); + std::optional generator = std::nullopt); + +std::shared_ptr Dropout(const std::shared_ptr &tensor, float p = 0.5f, + std::optional generator = std::nullopt); 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..aa429928 100644 --- a/infini_train/include/tensor.h +++ b/infini_train/include/tensor.h @@ -10,6 +10,7 @@ #include "Eigen/Dense" #include "glog/logging.h" +#include "infini_train/include/core/generator.h" #include "infini_train/include/datatype.h" #include "infini_train/include/device.h" #include "infini_train/include/scalar.h" @@ -151,7 +152,7 @@ class Tensor : public std::enable_shared_from_this { // distribution std::shared_ptr Uniform(float from = 0.0f, float to = 1.0f, - std::optional generator = std::nullopt); + std::optional generator = std::nullopt); std::shared_ptr Matmul(const std::shared_ptr &other); std::shared_ptr Outer(const std::shared_ptr &other); diff --git a/infini_train/src/core/cpu_generator.cc b/infini_train/src/core/cpu_generator.cc new file mode 100644 index 00000000..0f6dac89 --- /dev/null +++ b/infini_train/src/core/cpu_generator.cc @@ -0,0 +1,94 @@ +#include "infini_train/include/core/cpu_generator.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/device.h" + +namespace infini_train::core { +namespace { +// Backend magic prefixing every serialized CPU state. Used by SetState() to +// reject states produced by a different backend. +constexpr uint64_t kCPUStateMagic = 0x494E464350554731ULL; // "INFCPUG1" +// Explicit format version for forward-compatible state evolution. +constexpr uint64_t kCPUStateVersion = 1; + +using detail::AppendU64; +using detail::ReadU64; +} // namespace + +CPUGeneratorImpl::CPUGeneratorImpl(uint64_t seed) : seed_(seed), engine_(seed) {} + +void CPUGeneratorImpl::SetCurrentSeed(uint64_t seed) { + seed_ = seed; + engine_.seed(seed); +} + +uint64_t CPUGeneratorImpl::CurrentSeed() const { return seed_; } + +uint64_t CPUGeneratorImpl::Seed() { + std::random_device rd; + uint64_t seed = (static_cast(rd()) << 32) | static_cast(rd()); + SetCurrentSeed(seed); + return seed; +} + +void CPUGeneratorImpl::SetOffset(uint64_t offset) { + CHECK_EQ(offset, 0u) << "CPUGeneratorImpl does not support non-zero offset"; +} + +uint64_t CPUGeneratorImpl::GetOffset() const { return 0; } + +std::vector CPUGeneratorImpl::GetState() const { + // Layout: [magic u64][version u64][seed u64][engine-text-length u64][engine text bytes]. + std::ostringstream oss; + oss << engine_; + const std::string engine_text = oss.str(); + + std::vector state; + state.reserve(32 + engine_text.size()); + AppendU64(state, kCPUStateMagic); + AppendU64(state, kCPUStateVersion); + AppendU64(state, seed_); + AppendU64(state, static_cast(engine_text.size())); + state.insert(state.end(), engine_text.begin(), engine_text.end()); + return state; +} + +void CPUGeneratorImpl::SetState(const std::vector &state) { + CHECK_GE(state.size(), static_cast(32)) << "Invalid CPU generator state: too short"; + const uint64_t magic = ReadU64(state.data()); + CHECK_EQ(magic, kCPUStateMagic) << "Invalid CPU generator state: backend magic mismatch " + "(state may come from a non-CPU generator)"; + const uint64_t version = ReadU64(state.data() + 8); + CHECK_EQ(version, kCPUStateVersion) << "Invalid CPU generator state: unsupported version " << version + << " (expected " << kCPUStateVersion << ")"; + const uint64_t seed = ReadU64(state.data() + 16); + const uint64_t text_len = ReadU64(state.data() + 24); + CHECK_EQ(state.size(), static_cast(32) + text_len) << "Invalid CPU generator state: length mismatch"; + + const std::string engine_text(reinterpret_cast(state.data() + 32), text_len); + std::istringstream iss(engine_text); + iss >> engine_; + CHECK(!iss.fail()) << "Invalid CPU generator state: engine deserialization failed"; + seed_ = seed; +} + +Device CPUGeneratorImpl::GetDevice() const { return Device(Device::DeviceType::kCPU, 0); } + +std::shared_ptr CPUGeneratorImpl::Clone() const { + auto cloned = std::make_shared(seed_); + cloned->engine_ = engine_; + return cloned; +} + +std::mt19937_64 &CPUGeneratorImpl::Engine() { return engine_; } + +} // namespace infini_train::core diff --git a/infini_train/src/core/cuda_generator.cc b/infini_train/src/core/cuda_generator.cc new file mode 100644 index 00000000..cfe3bfce --- /dev/null +++ b/infini_train/src/core/cuda_generator.cc @@ -0,0 +1,100 @@ +#ifdef USE_CUDA + +#include "infini_train/include/core/cuda_generator.h" + +#include +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/device.h" + +namespace infini_train::core { +namespace { +constexpr uint64_t kCUDAStateMagic = 0x494E464355444731ULL; // "INFCUDG1" +// Explicit format version for forward-compatible state evolution. +constexpr uint64_t kCUDAStateVersion = 1; + +using detail::AppendU64; +using detail::ReadU64; +} // namespace + +CUDAGeneratorImpl::CUDAGeneratorImpl(int8_t device_index, uint64_t seed) + : device_index_(device_index), seed_(seed), philox_offset_(0) {} + +void CUDAGeneratorImpl::SetCurrentSeed(uint64_t seed) { + seed_ = seed; + philox_offset_ = 0; +} + +uint64_t CUDAGeneratorImpl::CurrentSeed() const { return seed_; } + +uint64_t CUDAGeneratorImpl::Seed() { + std::random_device rd; + // Limit to 53 bits to ensure unique representation when cast to double. + // This matches PyTorch's getNonDeterministicRandom(is_cuda=true) behavior. + uint64_t seed = ((static_cast(rd()) << 32) | static_cast(rd())) & 0x1FFFFFFFFFFFFF; + SetCurrentSeed(seed); + return seed; +} + +void CUDAGeneratorImpl::SetOffset(uint64_t offset) { philox_offset_ = offset; } + +uint64_t CUDAGeneratorImpl::GetOffset() const { return philox_offset_; } + +std::vector CUDAGeneratorImpl::GetState() const { + // Layout: [magic u64][version u64][seed u64][philox_offset u64][device_index u64]. + // device_index_ (int8_t) is widened to uint64_t for uniform 8-byte field + // alignment. The truncation in SetState is safe because device indices + // are always non-negative and well within [0, 127]. + std::vector state; + state.reserve(40); + AppendU64(state, kCUDAStateMagic); + AppendU64(state, kCUDAStateVersion); + AppendU64(state, seed_); + AppendU64(state, philox_offset_); + AppendU64(state, static_cast(device_index_)); + return state; +} + +void CUDAGeneratorImpl::SetState(const std::vector &state) { + CHECK_EQ(state.size(), static_cast(40)) << "Invalid CUDA generator state: unexpected length"; + const uint64_t magic = ReadU64(state.data()); + CHECK_EQ(magic, kCUDAStateMagic) << "Invalid CUDA generator state: backend magic mismatch " + "(state may come from a non-CUDA generator)"; + const uint64_t version = ReadU64(state.data() + 8); + CHECK_EQ(version, kCUDAStateVersion) << "Invalid CUDA generator state: unsupported version " << version + << " (expected " << kCUDAStateVersion << ")"; + // Truncation back to int8_t is safe: device indices are always in [0, 127]. + // The device_index check enforces that a state produced by a different CUDA + // device cannot be loaded into this generator, preventing unintended + // cross-device state aliasing (stricter than PyTorch, which omits this check). + const int8_t state_device_index = static_cast(ReadU64(state.data() + 32)); + CHECK_EQ(state_device_index, device_index_) + << "Invalid CUDA generator state: device index mismatch, expected " << static_cast(device_index_) + << " but found " << static_cast(state_device_index); + seed_ = ReadU64(state.data() + 16); + philox_offset_ = ReadU64(state.data() + 24); +} + +Device CUDAGeneratorImpl::GetDevice() const { return Device(Device::DeviceType::kCUDA, device_index_); } + +std::shared_ptr CUDAGeneratorImpl::Clone() const { + auto cloned = std::make_shared(device_index_, seed_); + cloned->philox_offset_ = philox_offset_; + return cloned; +} + +std::pair CUDAGeneratorImpl::PhiloxEngineInputs(uint64_t increment) { + const uint64_t offset = philox_offset_; + philox_offset_ += increment; + return {seed_, offset}; +} + +} // namespace infini_train::core + +#endif // USE_CUDA diff --git a/infini_train/src/core/generator.cc b/infini_train/src/core/generator.cc new file mode 100644 index 00000000..a02b2249 --- /dev/null +++ b/infini_train/src/core/generator.cc @@ -0,0 +1,103 @@ +#include "infini_train/include/core/generator.h" + +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/core/cpu_generator.h" +#include "infini_train/include/device.h" + +#ifdef USE_CUDA +#include "infini_train/include/core/cuda_generator.h" +#endif + +namespace infini_train::core { +namespace detail { +namespace { +// Guards the lazily-created default generators. +std::mutex &DefaultGeneratorMutex() { + static std::mutex mutex; + return mutex; +} + +std::optional &DefaultManualSeed() { + static std::optional seed; + return seed; +} + +// Process-wide default CPU generator (single shared random state source). +Generator &MutableDefaultCPUGenerator() { + static Generator generator(std::make_shared()); + return generator; +} + +#ifdef USE_CUDA +// Per-device default CUDA generators. Different devices are fully independent. +std::map &MutableDefaultCUDAGenerators() { + static std::map generators; + return generators; +} + +Generator &MutableDefaultCUDAGenerator(int8_t device_index) { + auto &generators = MutableDefaultCUDAGenerators(); + auto it = generators.find(device_index); + if (it == generators.end()) { + auto manual_seed = DefaultManualSeed(); + Generator generator = manual_seed.has_value() + ? Generator(std::make_shared(device_index, *manual_seed)) + : Generator(std::make_shared(device_index)); + it = generators.emplace(device_index, std::move(generator)).first; + } + return it->second; +} +#endif +} // namespace + +const Generator &DefaultCPUGenerator() { + std::lock_guard lock(DefaultGeneratorMutex()); + return MutableDefaultCPUGenerator(); +} + +const Generator &DefaultCUDAGenerator(int8_t device_index) { +#ifdef USE_CUDA + std::lock_guard lock(DefaultGeneratorMutex()); + return MutableDefaultCUDAGenerator(device_index); +#else + LOG(FATAL) << "CUDA default generator requested but the framework was built without CUDA support"; + return DefaultCPUGenerator(); +#endif +} + +const Generator &GetDefaultGenerator(const Device &device) { + switch (device.type()) { + case Device::DeviceType::kCPU: + return DefaultCPUGenerator(); + case Device::DeviceType::kCUDA: + return DefaultCUDAGenerator(device.index()); + default: + LOG(FATAL) << "No default generator for device type " << static_cast(device.type()); + return DefaultCPUGenerator(); + } +} +} // namespace detail + +void ManualSeed(uint64_t seed) { + // Lock ordering: DefaultGeneratorMutex (map guard) is always acquired first; + // Generator::SetCurrentSeed then acquires impl->mutex_ (impl guard). + // Callers must not hold impl->mutex_ when calling ManualSeed() to avoid + // potential deadlock with threads that draw from the default generator. + std::lock_guard lock(detail::DefaultGeneratorMutex()); + detail::DefaultManualSeed() = seed; + detail::MutableDefaultCPUGenerator().SetCurrentSeed(seed); +#ifdef USE_CUDA + // Seed already materialized CUDA generators. Future CUDA defaults read + // DefaultManualSeed() when they are lazily created. + for (auto &[index, generator] : detail::MutableDefaultCUDAGenerators()) { generator.SetCurrentSeed(seed); } +#endif +} + +} // namespace infini_train::core diff --git a/infini_train/src/kernels/cpu/distribution.cc b/infini_train/src/kernels/cpu/distribution.cc new file mode 100644 index 00000000..929ca4dc --- /dev/null +++ b/infini_train/src/kernels/cpu/distribution.cc @@ -0,0 +1,45 @@ +#include +#include +#include +#include +#include + +#include "infini_train/include/core/cpu_generator.h" +#include "infini_train/include/core/generator.h" +#include "infini_train/include/device.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +namespace { +template void FillKernel(float *data, int64_t numel, core::CPUGeneratorImpl *impl, Dist dist) { + std::lock_guard lock(impl->mutex_); + auto &engine = impl->Engine(); + for (int64_t i = 0; i < numel; ++i) { data[i] = dist(engine); } +} +} // namespace +// Convenience wrappers matching PyTorch's named distribution templates. + +// uniform_impl_: fills buffer with values from [from, to). +void Uniform(std::shared_ptr &tensor, float from, float to, core::CPUGeneratorImpl *generator) { + const int64_t numel = tensor->NumElements(); + float *buffer = static_cast(tensor->DataPtr()); + FillKernel(buffer, numel, generator, std::uniform_real_distribution(from, to)); +} + +// normal_impl_: fills buffer with values from N(mean, std). +void Normal(std::shared_ptr &tensor, float mean, float std, core::CPUGeneratorImpl *generator) { + float *buffer = static_cast(tensor->DataPtr()); + const int64_t numel = tensor->NumElements(); + FillKernel(buffer, numel, generator, std::normal_distribution(mean, std)); +} + +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_DISTRIBUTION_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_DISTRIBUTION_KERNEL(Uniform) +REGISTER_CPU_DISTRIBUTION_KERNEL(Normal) + +#undef REGISTER_CPU_DISTRIBUTION_KERNEL \ No newline at end of file diff --git a/infini_train/src/kernels/cpu/drop.cc b/infini_train/src/kernels/cpu/drop.cc new file mode 100644 index 00000000..c7a205f6 --- /dev/null +++ b/infini_train/src/kernels/cpu/drop.cc @@ -0,0 +1,39 @@ +#include +#include +#include +#include +#include + +#include "infini_train/include/core/cpu_generator.h" +#include "infini_train/include/core/generator.h" +#include "infini_train/include/device.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include +namespace infini_train::kernels::cpu { + +std::shared_ptr Dropout(const std::shared_ptr &input, float p, core::CPUGeneratorImpl *impl) { + CHECK(impl != nullptr); + auto device = input->GetDevice(); + float keep_prob = 1.0f - p; + float scale = 1.0f / keep_prob; + auto mask = std::make_shared(input->Dims(), DataType::kFLOAT32, device); + auto *input_ptr = static_cast(input->DataPtr()); + int64_t numel = mask->NumElements(); + std::bernoulli_distribution dist(keep_prob); + { + std::lock_guard lock(impl->mutex_); + auto &engine = impl->Engine(); + for (int64_t i = 0; i < numel; i++) { input_ptr[i] = dist(engine) ? input_ptr[i] * scale : 0.0f; } + } + return input; +} + +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_DROPOUT_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_DROPOUT_KERNEL(Dropout) +#undef REGISTER_CPU_DROPOUT_KERNEL \ No newline at end of file diff --git a/infini_train/src/kernels/cuda/common/policy_helper.cu b/infini_train/src/kernels/cuda/common/policy_helper.cu new file mode 100644 index 00000000..3356b253 --- /dev/null +++ b/infini_train/src/kernels/cuda/common/policy_helper.cu @@ -0,0 +1,32 @@ +#include "infini_train/src/kernels/cuda/common/policy_helper.cuh" + +#include + +namespace infini_train::kernels::cuda::common { + +LaunchConfig MakeElementwiseLaunch(int64_t numel, int block_size) { + + int device = 0; + cudaGetDevice(&device); + + cudaDeviceProp prop; + cudaGetDeviceProperties(&prop, device); + + int blocks_per_sm = prop.maxThreadsPerMultiProcessor / block_size; + + int max_blocks = prop.multiProcessorCount * blocks_per_sm; + + int grid = static_cast((numel + block_size - 1) / block_size); + + grid = std::min(grid, max_blocks); + + return {dim3(grid), dim3(block_size)}; +} +uint64_t CalcPhiloxCounterOffset(int64_t numel, int64_t threads, int unroll) { + + uint64_t iters_per_thread = (numel + threads * unroll - 1) / (threads * unroll); + + return iters_per_thread * unroll; +} + +} // namespace infini_train::kernels::cuda::common diff --git a/infini_train/src/kernels/cuda/common/policy_helper.cuh b/infini_train/src/kernels/cuda/common/policy_helper.cuh new file mode 100644 index 00000000..7af19ffd --- /dev/null +++ b/infini_train/src/kernels/cuda/common/policy_helper.cuh @@ -0,0 +1,17 @@ +#pragma once + +#include +#include + +namespace infini_train::kernels::cuda::common { + +struct LaunchConfig { + dim3 grid; + dim3 block; +}; + +LaunchConfig MakeElementwiseLaunch(int64_t numel, int block_size = 256); + +uint64_t CalcPhiloxCounterOffset(int64_t numel, int64_t threads, int unroll = 4); + +} // namespace infini_train::kernels::cuda::common \ No newline at end of file diff --git a/infini_train/src/kernels/cuda/distribution.cu b/infini_train/src/kernels/cuda/distribution.cu new file mode 100644 index 00000000..5f73b2c5 --- /dev/null +++ b/infini_train/src/kernels/cuda/distribution.cu @@ -0,0 +1,226 @@ +#include +#include +#include +#include +#include + +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/core/cuda_generator.h" +#include "infini_train/include/core/generator.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" +#include "infini_train/src/kernels/cuda/common/policy_helper.cuh" + +#include "infini_train/src/core/runtime/cuda/cuda_dispatch.h" +#include "infini_train/src/core/runtime/cuda/cuda_runtime_common.h" + +namespace infini_train::kernels::cuda { + +namespace { + +// --------------------------------------------------------------------------- +// Constants +// --------------------------------------------------------------------------- + +// Philox generates 128 bits = 4 × 32-bit floats per curand_uniform4 / +// curand_normal4 call. UNROLL must equal 4. +// (See cuda_DistributionTemplates.h — const int UNROLL = 4) +constexpr int kUnroll = 4; +constexpr int kBlockSize = 256; + +// --------------------------------------------------------------------------- +// CUDA kernels +// --------------------------------------------------------------------------- + +/** + * UniformPhiloxKernel + * + * Fills data[0..numel) with values drawn from Uniform[from, to). + * + * Design (mirrors distribution_elementwise_grid_stride_kernel in cuda_DistributionTemplates.h): + * + * - Each thread initialises its own curandStatePhilox4_32_10_t using the + * (seed, offset) pair supplied by the host. The unique per-thread + * subsequence index is the thread's global linear index `idx`, so all + * threads generate non-overlapping random sequences. + * + * - The grid-stride outer loop runs for `rounded_size / stride` iterations, + * where rounded_size is padded to a multiple of stride*kUnroll so every + * thread executes the same number of curand4 calls — a requirement for + * correct Philox offset accounting. + * + * - curand_uniform4 returns 4 floats in (0, 1]. We map to [from, to) by + * scaling and flipping the boundary (matches PyTorch's uniform_kernel). + */ +template +__global__ void UniformKernel(T *data, T from, T to, int64_t numel, uint64_t seed, uint64_t offset) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + + // Initialise per-thread Philox state. + // curand_init(seed, sequence=idx, offset, &state): + // - `sequence` selects an independent 2^67-element subsequence per thread. + // - `offset` skips the first `offset` values within that subsequence, + // allowing successive kernel launches to continue from where the + // previous one left off without overlap. + curandStatePhilox4_32_10_t state; + curand_init(seed, static_cast(idx), offset, &state); + + const float range = to - from; + const int64_t stride = static_cast(blockDim.x) * gridDim.x; + + // Pad to a multiple of stride*kUnroll so every thread runs the same number + // of outer iterations (and thus the same number of curand4 calls). + const int64_t rounded_size = ((numel - 1) / (stride * kUnroll) + 1) * stride * kUnroll; + + for (int64_t linear_index = idx; linear_index < rounded_size; linear_index += stride * kUnroll) { + // curand_uniform4 advances the Philox counter by kUnroll=4 positions + // and returns 4 floats in (0, 1]. + float4 rand = curand_uniform4(&state); + +#pragma unroll + for (int ii = 0; ii < kUnroll; ++ii) { + const int64_t li = linear_index + stride * ii; + if (li < numel) { + const float r = (&rand.x)[ii]; + // Map (0, 1] → [from, to). If the scaled value hits exactly + // `to` (rare due to floating-point rounding), clamp to `from` + // to preserve the half-open interval semantics. + // (See cuda_DistributionTemplates.h uniform_kernel.) + float val = r * range + from; + val = (val == to) ? from : val; + data[li] = val; + } + } + } +} + +/** + * NormalPhiloxKernel + * + * Fills data[0..numel) with values drawn from N(mean, std_val). + * + * Uses curand_normal4, which applies the Box-Muller transform internally to + * produce 4 standard-normal floats per call. We then scale by std_val and + * shift by mean. + * + * (See cuda_DistributionTemplates.h normal_and_transform / + * normal_kernel for the PyTorch equivalent.) + */ +template +__global__ void NormalKernel(T *data, T mean, T std_val, int64_t numel, uint64_t seed, uint64_t offset) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + + curandStatePhilox4_32_10_t state; + curand_init(seed, static_cast(idx), offset, &state); + + const int64_t stride = static_cast(blockDim.x) * gridDim.x; + const int64_t rounded_size = ((numel - 1) / (stride * kUnroll) + 1) * stride * kUnroll; + + for (int64_t linear_index = idx; linear_index < rounded_size; linear_index += stride * kUnroll) { + // curand_normal4 returns 4 standard-normal floats (mean=0, std=1). + float4 rand = curand_normal4(&state); + +#pragma unroll + for (int ii = 0; ii < kUnroll; ++ii) { + const int64_t li = linear_index + stride * ii; + if (li < numel) { + data[li] = (&rand.x)[ii] * std_val + mean; + } + } + } +} + +} // namespace + +// --------------------------------------------------------------------------- +// Host-side launcher implementations +// --------------------------------------------------------------------------- + +void Uniform(std::shared_ptr tensor, float mean, float std_val, core::CUDAGeneratorImpl *impl) { + + CHECK(impl != nullptr) << "Generator must be a CUDAGeneratorImpl for CUDA kernels"; + CHECK(tensor != nullptr) << "Tensor must not be null"; + const int64_t numel = tensor->NumElements(); + if (numel == 0) { + return; + } + + auto launch = common::MakeElementwiseLaunch(numel, kBlockSize); + + uint64_t counter_offset + = common::CalcPhiloxCounterOffset(numel, static_cast(launch.grid.x) * launch.block.x, kUnroll); + + const auto &cuda_stream + = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(tensor->GetDevice().type())->GetStream(tensor->GetDevice())) + ->cuda_stream(); + // Acquire the generator lock only for the brief PhiloxEngineInputs call, + // then release before launching the kernel. The kernel runs purely on + // local (seed, offset) values; no device-side synchronisation is needed. + // (See "Note [Acquire lock when using random generators]" in PyTorch.) + uint64_t seed = 0; + uint64_t offset = 0; + { + std::lock_guard lock(impl->mutex_); + std::tie(seed, offset) = impl->PhiloxEngineInputs(counter_offset); + } + // core::cuda::DispatchCudaFunc( + // tensor->Dtype(), + // [=]() { + // const T casted_mean = mean.to(); + // const T casted_std = std_val.to(); + UniformKernel<<>>(static_cast(tensor->DataPtr()), mean, + std_val, numel, seed, offset); + // } + // ); +} + +void Normal(std::shared_ptr tensor, float mean, float std_val, core::CUDAGeneratorImpl *impl) { + + CHECK(impl != nullptr) << "Generator must be a CUDAGeneratorImpl for CUDA kernels"; + CHECK(tensor != nullptr) << "Tensor must not be null"; + const int64_t numel = tensor->NumElements(); + if (numel == 0) { + return; + } + + auto launch = common::MakeElementwiseLaunch(numel, kBlockSize); + + uint64_t counter_offset + = common::CalcPhiloxCounterOffset(numel, static_cast(launch.grid.x) * launch.block.x, kUnroll); + + const auto &cuda_stream + = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(tensor->GetDevice().type())->GetStream(tensor->GetDevice())) + ->cuda_stream(); + uint64_t seed = 0; + uint64_t offset = 0; + { + std::lock_guard lock(impl->mutex_); + std::tie(seed, offset) = impl->PhiloxEngineInputs(counter_offset); + } + // core::cuda::DispatchCudaFunc( + // tensor->Dtype(), + // [=]() { + // const T casted_mean = static_cast(mean); + // const T casted_std = static_cast(std_val); + NormalKernel<<>>(static_cast(tensor->DataPtr()), mean, + std_val, numel, seed, offset); + // } + // ); +} + +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_DISTRIBUTION_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_DISTRIBUTION_KERNEL(Uniform) +REGISTER_CUDA_DISTRIBUTION_KERNEL(Normal) + +#undef REGISTER_CUDA_DISTRIBUTION_KERNEL \ No newline at end of file diff --git a/infini_train/src/kernels/cuda/drop.cu b/infini_train/src/kernels/cuda/drop.cu new file mode 100644 index 00000000..0e02b501 --- /dev/null +++ b/infini_train/src/kernels/cuda/drop.cu @@ -0,0 +1,96 @@ +#include +#include +#include +#include + +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/core/cuda_generator.h" +#include "infini_train/include/core/generator.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/cuda/cuda_runtime_common.h" +#include "infini_train/src/kernels/cuda/common/policy_helper.cuh" +namespace infini_train::kernels::cuda { + +namespace { + +constexpr int kUnroll = 4; +constexpr int kBlockSize = 256; + +__global__ void DropoutKernel(float *data, int64_t numel, uint64_t seed, uint64_t offset, float p) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + + curandStatePhilox4_32_10_t state; + curand_init(seed, static_cast(idx), offset, &state); + + const int64_t stride = static_cast(blockDim.x) * gridDim.x; + const int64_t rounded_size = ((numel - 1) / (stride * kUnroll) + 1) * stride * kUnroll; + + const float keep_scale = 1.0f / (1.0f - p); + + for (int64_t linear_index = idx; linear_index < rounded_size; linear_index += stride * kUnroll) { + float4 rand = curand_uniform4(&state); + +#pragma unroll + for (int ii = 0; ii < kUnroll; ++ii) { + const int64_t li = linear_index + stride * ii; + if (li < numel) { + const float r = (&rand.x)[ii]; + data[li] = (r < p) ? 0.0f : data[li] * keep_scale; + } + } + } +} + +} // namespace + +std::shared_ptr Dropout(std::shared_ptr tensor, float p, core::CUDAGeneratorImpl *impl) { + auto device = tensor->GetDevice(); + core::DeviceGuard guard(device); + + CHECK(device.IsCUDA()) << "CUDA Dropout kernel requires a CUDA tensor"; + CHECK_EQ(static_cast(tensor->Dtype()), static_cast(DataType::kFLOAT32)) + << "CUDA Dropout currently only supports FLOAT32 tensors"; + + const int64_t numel = tensor->NumElements(); + if (numel == 0) { + return tensor; + } + + auto launch = common::MakeElementwiseLaunch(numel, kBlockSize); + + uint64_t counter_offset + = common::CalcPhiloxCounterOffset(numel, static_cast(launch.grid.x) * launch.block.x, kUnroll); + + const auto &cuda_stream + = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(tensor->GetDevice().type())->GetStream(tensor->GetDevice())) + ->cuda_stream(); + + uint64_t seed = 0; + uint64_t offset = 0; + { + std::lock_guard lock(impl->mutex_); + std::tie(seed, offset) = impl->PhiloxEngineInputs(counter_offset); + } + + DropoutKernel<<>>(static_cast(tensor->DataPtr()), numel, + seed, offset, p); + + return tensor; +} + +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_DROPOUT_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_DROPOUT_KERNEL(Dropout) + +#undef REGISTER_CUDA_DROPOUT_KERNEL \ No newline at end of file diff --git a/infini_train/src/nn/init.cc b/infini_train/src/nn/init.cc index 79b4b48b..bca8788d 100644 --- a/infini_train/src/nn/init.cc +++ b/infini_train/src/nn/init.cc @@ -7,62 +7,37 @@ #include #include -#ifdef USE_OMP -#include -#endif - #include "glog/logging.h" +#include "infini_train/include/core/cpu_generator.h" +#include "infini_train/include/core/cuda_generator.h" +#include "infini_train/include/core/generator.h" #include "infini_train/include/core/runtime/device_guard.h" #include "infini_train/include/device.h" +#include "infini_train/include/dispatcher.h" #include "infini_train/include/tensor.h" +#include 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); -} // namespace std::shared_ptr Normal(const std::shared_ptr &tensor, float mean, float std, - std::optional generator) { - 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); - } - } -#else - std::normal_distribution dis(mean, std); - std::generate(buffer.begin(), buffer.end(), [&]() { return generator ? dis(generator.value()) : dis(gen); }); -#endif - + std::optional generator) { + // TODO(dcj): Support other floating point data types later. + CHECK_EQ(static_cast(tensor->Dtype()), static_cast(DataType::kFLOAT32)) + << "Normal initialization currently only supports FLOAT32 tensors"; auto device = tensor->GetDevice(); - core::DeviceGuard guard(device); - auto impl = core::GetDeviceGuardImpl(device.type()); - - impl->MemcpyAsync(tensor->DataPtr(), buffer.data(), num_elements * sizeof(float), - device.type() == Device::DeviceType::kCPU ? core::MemcpyKind::kD2D : core::MemcpyKind::kH2D, - impl->GetStream(device)); - return tensor; + auto kernel = Dispatcher::Instance().GetKernel({device.type(), "Normal"}); + if (device.IsCPU()) { + auto gen = core::GetGeneratorOrDefault(generator, device); + kernel.Call(tensor, mean, std, gen); + return tensor; + } else if (device.IsCUDA()) { + auto gen = core::GetGeneratorOrDefault(generator, device); + kernel.Call(tensor, mean, std, gen); + return tensor; + } else { + LOG(FATAL) << "Unsupported device type: " << static_cast(device.type()); + } } std::pair CalculateFanInAndFanOut(const std::shared_ptr &tensor) { @@ -113,7 +88,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::optional generator) { for (const auto dim : tensor->Dims()) { if (dim == 0) { LOG(WARNING) << "Initializing zero-element tensors is a no-op"; @@ -128,35 +103,44 @@ 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) { - 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); - } - } -#else - std::uniform_real_distribution dis(a, b); - std::generate(buffer.begin(), buffer.end(), [&]() { return generator ? dis(generator.value()) : dis(gen); }); -#endif - + std::optional generator) { + // TODO(dcj): Support other floating point data types later. + CHECK_EQ(static_cast(tensor->Dtype()), static_cast(DataType::kFLOAT32)) + << "Uniform initialization currently only supports FLOAT32 tensors"; auto device = tensor->GetDevice(); + auto kernel = Dispatcher::Instance().GetKernel({device.type(), "Uniform"}); + if (device.IsCPU()) { + auto gen = core::GetGeneratorOrDefault(generator, device); + kernel.Call(tensor, a, b, gen); + return tensor; + } else if (device.IsCUDA()) { + auto gen = core::GetGeneratorOrDefault(generator, device); + kernel.Call(tensor, a, b, gen); + return tensor; + } else { + LOG(FATAL) << "Unsupported device type: " << static_cast(device.type()); + } +} - core::DeviceGuard guard(device); - auto impl = core::GetDeviceGuardImpl(device.type()); - - impl->MemcpyAsync(tensor->DataPtr(), buffer.data(), num_elements * sizeof(float), - device.type() == Device::DeviceType::kCPU ? core::MemcpyKind::kD2D : core::MemcpyKind::kH2D, - impl->GetStream(device)); +std::shared_ptr Dropout(const std::shared_ptr &tensor, float p, + std::optional generator) { + CHECK_GE(p, 0.0f) << "Dropout probability must be in [0, 1)"; + CHECK_LT(p, 1.0f) << "Dropout probability must be in [0, 1)"; + CHECK_EQ(static_cast(tensor->Dtype()), static_cast(DataType::kFLOAT32)) + << "Dropout currently only supports FLOAT32 tensors"; - return tensor; + auto device = tensor->GetDevice(); + auto kernel = Dispatcher::Instance().GetKernel({device.type(), "Dropout"}); + + if (device.IsCPU()) { + auto gen = core::GetGeneratorOrDefault(generator, device); + return kernel.Call>(tensor, p, gen); + } else if (device.IsCUDA()) { + auto gen = core::GetGeneratorOrDefault(generator, device); + return kernel.Call>(tensor, p, gen); + } else { + LOG(FATAL) << "Unsupported device type"; + } } std::shared_ptr Ones(const std::shared_ptr &tensor) { diff --git a/infini_train/src/tensor.cc b/infini_train/src/tensor.cc index 18ca3d22..5f72add3 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::optional generator) { return nn::init::Uniform(shared_from_this(), from, to, generator); } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 96776585..2971de1b 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -30,3 +30,6 @@ add_subdirectory(transformer) # Checkpoint tests add_subdirectory(checkpoint) + +# Generator (random number generator) tests +add_subdirectory(generator) diff --git a/tests/generator/CMakeLists.txt b/tests/generator/CMakeLists.txt new file mode 100644 index 00000000..0bef53ff --- /dev/null +++ b/tests/generator/CMakeLists.txt @@ -0,0 +1,10 @@ +# ========================================================================== +# Generator (random number generator) tests +# ========================================================================== + +file(GLOB GENERATOR_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/test_generator_*.cc) + +infini_train_add_test_suite(test_generator + SOURCES ${GENERATOR_SOURCES} +) + diff --git a/tests/generator/test_generator_api.cc b/tests/generator/test_generator_api.cc new file mode 100644 index 00000000..0b31db99 --- /dev/null +++ b/tests/generator/test_generator_api.cc @@ -0,0 +1,184 @@ +// Generator handle class API tests. +// +// These tests verify the basic semantics of the Generator wrapper/handle class +// (e.g. Defined(), comparison operators, copy/move semantics, and polymorphic +// dispatch using a simple stub implementation) without depending on the concrete +// CPU or CUDA generator backends. + +#include +#include + +#include "infini_train/include/core/generator.h" +#include "infini_train/include/device.h" + +#include "gtest/gtest.h" + +using namespace infini_train; + +namespace { + +// A simple stub generator implementation for verifying handle dispatch. +class StubGeneratorImpl : public core::GeneratorImpl { +public: + static constexpr Device::DeviceType kDeviceType = Device::DeviceType::kCPU; + + explicit StubGeneratorImpl(uint64_t seed = 42) : seed_(seed), offset_(0) {} + + void SetCurrentSeed(uint64_t seed) override { seed_ = seed; } + + uint64_t CurrentSeed() const override { return seed_; } + + uint64_t Seed() override { return ++seed_; } + + void SetOffset(uint64_t offset) override { offset_ = offset; } + + uint64_t GetOffset() const override { return offset_; } + + std::vector GetState() const override { return std::vector{0xBE, 0xEF}; } + + void SetState(const std::vector &state) override { (void)state; } + + Device GetDevice() const override { return Device(Device::DeviceType::kCPU, 0); } + + std::shared_ptr Clone() const override { + auto cloned = std::make_shared(seed_); + cloned->offset_ = offset_; + return cloned; + } + +private: + uint64_t seed_; + uint64_t offset_; +}; + +} // namespace + +TEST(GeneratorAPITest, DefaultConstructedIsUndefined) { + core::Generator gen; + EXPECT_FALSE(gen.Defined()); +} + +TEST(GeneratorAPITest, ConstructorRejectsNullptr) { + EXPECT_DEATH(core::Generator(nullptr), "nullptr is not supported"); +} + +TEST(GeneratorAPITest, DefinedWhenImplProvided) { + auto impl = std::make_shared(123); + core::Generator gen(impl); + EXPECT_TRUE(gen.Defined()); + EXPECT_EQ(gen.CurrentSeed(), 123u); +} + +TEST(GeneratorAPITest, CopySharesImplAndState) { + auto impl = std::make_shared(10); + core::Generator gen1(impl); + core::Generator gen2 = gen1; + + EXPECT_TRUE(gen1.Defined()); + EXPECT_TRUE(gen2.Defined()); + EXPECT_EQ(gen1, gen2); + EXPECT_EQ(gen1.UnsafeGetImpl(), gen2.UnsafeGetImpl()); + + // Modifying via one handle affects the other since they share the impl. + gen1.SetCurrentSeed(20); + EXPECT_EQ(gen2.CurrentSeed(), 20u); +} + +TEST(GeneratorAPITest, MoveTransfersImpl) { + auto impl = std::make_shared(10); + core::Generator gen1(impl); + core::Generator gen2 = std::move(gen1); + + EXPECT_FALSE(gen1.Defined()); + EXPECT_TRUE(gen2.Defined()); + EXPECT_EQ(gen2.CurrentSeed(), 10u); +} + +TEST(GeneratorAPITest, CloneCreatesIndependentImpl) { + auto impl = std::make_shared(10); + core::Generator gen1(impl); + core::Generator gen2 = gen1.Clone(); + + EXPECT_TRUE(gen1.Defined()); + EXPECT_TRUE(gen2.Defined()); + EXPECT_NE(gen1, gen2); + EXPECT_NE(gen1.UnsafeGetImpl(), gen2.UnsafeGetImpl()); + + EXPECT_EQ(gen1.CurrentSeed(), gen2.CurrentSeed()); + + // Mutating gen1 does not affect gen2 because they are clones. + gen1.SetCurrentSeed(30); + EXPECT_EQ(gen1.CurrentSeed(), 30u); + EXPECT_EQ(gen2.CurrentSeed(), 10u); +} + +TEST(GeneratorAPITest, BasicAPIForwarding) { + auto impl = std::make_shared(100); + core::Generator gen(impl); + + EXPECT_EQ(gen.CurrentSeed(), 100u); + + gen.ManualSeed(200); + EXPECT_EQ(gen.CurrentSeed(), 200u); + + uint64_t old_seed = gen.CurrentSeed(); + uint64_t new_seed = gen.Seed(); + EXPECT_NE(new_seed, old_seed); + EXPECT_EQ(new_seed, gen.CurrentSeed()); + + EXPECT_EQ(gen.GetOffset(), 0u); + gen.SetOffset(99); + EXPECT_EQ(gen.GetOffset(), 99u); + + std::vector expected_state{0xBE, 0xEF}; + EXPECT_EQ(gen.GetState(), expected_state); + + EXPECT_EQ(gen.GetDevice(), Device(Device::DeviceType::kCPU, 0)); +} + +TEST(GeneratorAPITest, CheckGeneratorDeath) { + std::optional null_gen = std::nullopt; + EXPECT_DEATH(core::CheckGenerator(null_gen), "Expected a Generator but received std::nullopt"); + + core::Generator undef_gen; + EXPECT_DEATH(core::CheckGenerator(undef_gen), "undefined implementation"); +} + +TEST(GeneratorAPITest, MakeGeneratorTest) { + auto gen = core::MakeGenerator(999); + EXPECT_TRUE(gen.Defined()); + EXPECT_EQ(gen.CurrentSeed(), 999u); +} + +TEST(GeneratorAPITest, UnsafeGetImplBehavior) { + auto gen = core::MakeGenerator(111); + auto *impl = gen.UnsafeGetImpl(); + EXPECT_NE(impl, nullptr); + EXPECT_EQ(impl->CurrentSeed(), 111u); + + // Modifying via UnsafeGetImpl directly changes the generator state + impl->SetCurrentSeed(222); + EXPECT_EQ(gen.CurrentSeed(), 222u); +} + +TEST(GeneratorAPITest, CurrentSeedRelationship) { + auto gen = core::MakeGenerator(100); + EXPECT_EQ(gen.CurrentSeed(), 100u); + + uint64_t old_seed = gen.CurrentSeed(); + uint64_t new_seed = gen.Seed(); + EXPECT_NE(new_seed, old_seed); + EXPECT_EQ(new_seed, gen.CurrentSeed()); +} + +TEST(GeneratorAPITest, EqualityOperator) { + auto gen1 = core::MakeGenerator(1); + + auto gen2 = gen1; + + auto gen3 = gen1.Clone(); + + EXPECT_TRUE(gen1 == gen2); + + EXPECT_FALSE(gen1 == gen3); +} diff --git a/tests/generator/test_generator_cpu.cc b/tests/generator/test_generator_cpu.cc new file mode 100644 index 00000000..779922d2 --- /dev/null +++ b/tests/generator/test_generator_cpu.cc @@ -0,0 +1,178 @@ +// CPU Generator unit tests. + +#include "infini_train/include/core/cpu_generator.h" +#include "infini_train/include/core/generator.h" +#include "infini_train/include/device.h" +#include "gtest/gtest.h" +#include +#include +#include +#include + +using namespace infini_train; + +namespace { + +core::Generator MakeCPUGenerator(uint64_t seed) { + return core::Generator(std::make_shared(seed)); +} + +std::vector DrawCPU(core::Generator &generator, int64_t count) { + auto *impl = generator.Get(); + std::lock_guard lock(generator.Mutex()); + std::vector values(count); + auto &engine = impl->Engine(); + for (auto &value : values) { value = engine(); } + return values; +} + +struct FakeCUDAGeneratorImpl : public core::GeneratorImpl { + static constexpr Device::DeviceType kDeviceType = Device::DeviceType::kCUDA; + void SetCurrentSeed(uint64_t) override {} + uint64_t CurrentSeed() const override { return 0; } + uint64_t Seed() override { return 0; } + void SetOffset(uint64_t) override {} + uint64_t GetOffset() const override { return 0; } + std::vector GetState() const override { return {}; } + void SetState(const std::vector &) override {} + Device GetDevice() const override { return Device(Device::DeviceType::kCUDA, 0); } + std::shared_ptr Clone() const override { return nullptr; } +}; + +} // namespace + +TEST(CPUGeneratorTest, SeedSetAndGet) { + auto gen = MakeCPUGenerator(123); + EXPECT_EQ(gen.CurrentSeed(), 123u); + + gen.SetCurrentSeed(456); + EXPECT_EQ(gen.CurrentSeed(), 456u); + + gen.ManualSeed(789); + EXPECT_EQ(gen.CurrentSeed(), 789u); + + const uint64_t drawn = gen.Seed(); + EXPECT_EQ(gen.CurrentSeed(), drawn); +} + +TEST(CPUGeneratorTest, DeviceType) { + auto gen = MakeCPUGenerator(1); + EXPECT_TRUE(gen.GetDevice().IsCPU()); + EXPECT_EQ(gen.GetDevice().type(), Device::DeviceType::kCPU); +} + +TEST(CPUGeneratorTest, CloneIsIndependentWithSameState) { + auto gen = MakeCPUGenerator(42); + auto prefix = DrawCPU(gen, 4); + EXPECT_EQ(prefix.size(), 4u); + + auto cloned = gen.Clone(); + EXPECT_NE(gen.UnsafeGetImpl(), cloned.UnsafeGetImpl()); + EXPECT_EQ(DrawCPU(gen, 8), DrawCPU(cloned, 8)); +} + +TEST(CPUGeneratorTest, SameSeedSameSequence) { + auto g1 = MakeCPUGenerator(2024); + auto g2 = MakeCPUGenerator(2024); + EXPECT_EQ(DrawCPU(g1, 64), DrawCPU(g2, 64)); +} + +TEST(CPUGeneratorTest, DifferentSeedDifferentSequence) { + auto g1 = MakeCPUGenerator(1); + auto g2 = MakeCPUGenerator(2); + EXPECT_NE(DrawCPU(g1, 64), DrawCPU(g2, 64)); +} + +TEST(CPUGeneratorTest, GetSetStateRestoresCPUSequence) { + auto gen = MakeCPUGenerator(99); + const auto state = gen.GetState(); + + const auto first = DrawCPU(gen, 32); + + gen.SetState(state); + const auto restored = DrawCPU(gen, 32); + + EXPECT_EQ(first, restored); +} + +TEST(CPUGeneratorTest, StateRoundTripPreservesSeed) { + auto gen = MakeCPUGenerator(0xABCDEF); + const auto state = gen.GetState(); + + auto other = MakeCPUGenerator(1); + other.SetState(state); + EXPECT_EQ(other.CurrentSeed(), 0xABCDEFu); +} + +TEST(CPUGeneratorTest, RejectsStateTooShort) { + auto gen = MakeCPUGenerator(1); + auto state = gen.GetState(); + state.resize(state.size() / 2); + EXPECT_DEATH(gen.SetState(state), "too short|Invalid CPU generator state"); +} + +TEST(CPUGeneratorTest, RejectsForeignBackendState) { + auto gen = MakeCPUGenerator(1); + std::vector foreign(gen.GetState().size(), 0); + EXPECT_DEATH(gen.SetState(foreign), "backend magic mismatch"); +} + +TEST(CPUGeneratorTest, RejectsUnsupportedCPUStateVersion) { + auto gen = MakeCPUGenerator(42); + auto state = gen.GetState(); + + // Corrupt version field (bytes 8-15) to 99. + state[8] = 99; + for (int i = 9; i < 16; ++i) { state[i] = 0; } + + EXPECT_DEATH(gen.SetState(state), "unsupported version"); +} + +TEST(CPUGeneratorTest, CPUOnlySupportsZeroOffset) { + auto gen = MakeCPUGenerator(7); + EXPECT_EQ(gen.GetOffset(), 0u); + gen.SetOffset(0); + EXPECT_EQ(gen.GetOffset(), 0u); + EXPECT_DEATH(gen.SetOffset(1), "non-zero offset"); +} + +TEST(CPUGeneratorTest, GetWrongTypeDeath) { + auto gen = MakeCPUGenerator(123); + EXPECT_DEATH(core::CheckGenerator(gen), "Generator device type mismatch"); + EXPECT_DEATH(core::GetGeneratorOrDefault(gen, Device(Device::DeviceType::kCPU, 0)), + "Generator device type mismatch"); +} + +TEST(CPUGeneratorTest, RejectsMalformedCPUState) { + auto gen = MakeCPUGenerator(1); + EXPECT_DEATH(gen.SetState(std::vector{1, 2, 3}), "Invalid CPU generator state"); +} + +TEST(CPUGeneratorTest, RejectsStateLengthMismatchTruncated) { + auto gen = MakeCPUGenerator(1); + auto state = gen.GetState(); + state.pop_back(); + EXPECT_DEATH(gen.SetState(state), "length mismatch"); +} + +TEST(CPUGeneratorTest, RejectsStateLengthMismatchExtra) { + auto gen = MakeCPUGenerator(1); + auto state = gen.GetState(); + state.push_back(0); + EXPECT_DEATH(gen.SetState(state), "length mismatch"); +} + +TEST(CPUGeneratorTest, RejectsInvalidMagic) { + auto gen = MakeCPUGenerator(1); + auto state = gen.GetState(); + state[0] ^= 0xFF; + EXPECT_DEATH(gen.SetState(state), "backend magic mismatch"); +} + +TEST(CPUGeneratorTest, RejectsGarbageEngineData) { + auto gen = MakeCPUGenerator(1); + auto state = gen.GetState(); + size_t data_start = state.size() / 2; + for (size_t i = data_start; i < state.size(); ++i) { state[i] = 0xFF; } + EXPECT_DEATH(gen.SetState(state), "engine deserialization failed"); +} diff --git a/tests/generator/test_generator_cuda.cc b/tests/generator/test_generator_cuda.cc new file mode 100644 index 00000000..f5136bab --- /dev/null +++ b/tests/generator/test_generator_cuda.cc @@ -0,0 +1,167 @@ +// CUDA Generator unit tests. + +#include "infini_train/include/core/generator.h" +#include "infini_train/include/device.h" +#include "gtest/gtest.h" +#include +#include +#include +#include + +#ifdef USE_CUDA +#include "infini_train/include/core/cuda_generator.h" +#endif + +using namespace infini_train; + +#ifdef USE_CUDA + +TEST(CUDAGeneratorTest, OffsetAdvancesAndRestores) { + core::Generator gen(std::make_shared(0, 123)); + auto *impl = gen.Get(); + + auto first = impl->PhiloxEngineInputs(10); + EXPECT_EQ(first.first, 123u); + EXPECT_EQ(first.second, 0u); + EXPECT_EQ(gen.GetOffset(), 10u); + + const auto state = gen.GetState(); + auto second = impl->PhiloxEngineInputs(5); + EXPECT_EQ(second.second, 10u); + + gen.SetState(state); + EXPECT_EQ(gen.GetOffset(), 10u); +} + +TEST(CUDAGeneratorTest, SetCurrentSeedResetsOffset) { + core::Generator gen(std::make_shared(0, 123)); + auto *impl = gen.Get(); + + { + std::lock_guard lock(impl->mutex_); + impl->PhiloxEngineInputs(32); + } + + EXPECT_EQ(gen.GetOffset(), 32u); + + gen.SetCurrentSeed(456); + EXPECT_EQ(gen.CurrentSeed(), 456u); + EXPECT_EQ(gen.GetOffset(), 0u); +} + +TEST(CUDAGeneratorTest, RejectsDifferentDeviceState) { + core::Generator gen0(std::make_shared(0, 123)); + core::Generator gen1(std::make_shared(1, 456)); + + const auto state_from_device_one = gen1.GetState(); + EXPECT_DEATH(gen0.SetState(state_from_device_one), "device index mismatch"); +} + +TEST(CUDAGeneratorTest, RejectsUnsupportedCUDAStateVersion) { + core::Generator gen(std::make_shared(0, 42)); + auto state = gen.GetState(); + + // Corrupt version field (bytes 8-15) to 99. + state[8] = 99; + for (int i = 9; i < 16; ++i) { state[i] = 0; } + + EXPECT_DEATH(gen.SetState(state), "unsupported version"); +} + +TEST(CUDAGeneratorTest, CloneIsIndependent) { + core::Generator gen(std::make_shared(0, 42)); + auto *impl = gen.Get(); + { + std::lock_guard lock(impl->mutex_); + impl->PhiloxEngineInputs(10); + } + + auto cloned = gen.Clone(); + EXPECT_NE(gen.UnsafeGetImpl(), cloned.UnsafeGetImpl()); + EXPECT_EQ(gen.GetOffset(), cloned.GetOffset()); + EXPECT_EQ(gen.CurrentSeed(), cloned.CurrentSeed()); + + // Advance original — clone must not follow. + { + std::lock_guard lock(impl->mutex_); + impl->PhiloxEngineInputs(5); + } + EXPECT_EQ(gen.GetOffset(), 15u); + EXPECT_EQ(cloned.GetOffset(), 10u); +} + +TEST(CUDAGeneratorTest, SetStateRoundTripPreservesSeed) { + core::Generator gen(std::make_shared(0, 0xCAFE)); + auto *impl = gen.Get(); + { + std::lock_guard lock(impl->mutex_); + impl->PhiloxEngineInputs(100); + } + const auto state = gen.GetState(); + + // Mutate, then restore. + gen.SetCurrentSeed(999); + EXPECT_EQ(gen.CurrentSeed(), 999u); + EXPECT_EQ(gen.GetOffset(), 0u); + + gen.SetState(state); + EXPECT_EQ(gen.CurrentSeed(), 0xCAFEu); + EXPECT_EQ(gen.GetOffset(), 100u); +} + +TEST(CUDAGeneratorTest, ZeroIncrementDoesNotAdvanceOffset) { + core::Generator gen(std::make_shared(0, 1)); + + auto *impl = gen.Get(); + + { + std::lock_guard lock(impl->mutex_); + auto state = impl->PhiloxEngineInputs(0); + + EXPECT_EQ(state.first, 1u); + EXPECT_EQ(state.second, 0u); + } + + EXPECT_EQ(gen.GetOffset(), 0u); +} + +TEST(CUDAGeneratorTest, MultipleSmallIncrements) { + core::Generator gen(std::make_shared(0, 1)); + + auto *impl = gen.Get(); + + { + std::lock_guard lock(impl->mutex_); + + impl->PhiloxEngineInputs(3); + impl->PhiloxEngineInputs(5); + impl->PhiloxEngineInputs(7); + } + + EXPECT_EQ(gen.GetOffset(), 15u); +} + +TEST(CUDAGeneratorTest, LargeIncrement) { + core::Generator gen(std::make_shared(0, 1)); + + auto *impl = gen.Get(); + + constexpr uint64_t kIncrement = 1ull << 20; + + { + std::lock_guard lock(impl->mutex_); + impl->PhiloxEngineInputs(kIncrement); + } + + EXPECT_EQ(gen.GetOffset(), kIncrement); +} + +TEST(CUDAGeneratorTest, SetOffsetWorks) { + core::Generator gen(std::make_shared(0, 1)); + + gen.SetOffset(1234); + + EXPECT_EQ(gen.GetOffset(), 1234u); +} + +#endif // USE_CUDA diff --git a/tests/generator/test_generator_default.cc b/tests/generator/test_generator_default.cc new file mode 100644 index 00000000..ae77a1fa --- /dev/null +++ b/tests/generator/test_generator_default.cc @@ -0,0 +1,141 @@ +// Default generator management and ManualSeed tests. + +#include "infini_train/include/core/generator.h" +#include "infini_train/include/device.h" +#include "gtest/gtest.h" +#include +#include +#include +#include + +#include "infini_train/include/core/cpu_generator.h" +#ifdef USE_CUDA +#include "infini_train/include/core/cuda_generator.h" +#endif + +using namespace infini_train; + +namespace { + +Device CPUDevice() { return Device(Device::DeviceType::kCPU, 0); } + +core::Generator MakeCPUGenerator(uint64_t seed) { + return core::Generator(std::make_shared(seed)); +} + +std::vector DrawCPU(core::Generator &generator, int64_t count) { + auto *impl = generator.Get(); + std::lock_guard lock(generator.Mutex()); + std::vector values(count); + auto &engine = impl->Engine(); + for (auto &value : values) { value = engine(); } + return values; +} + +} // namespace + +TEST(GeneratorDefaultTest, DefaultGeneratorIsStable) { + const auto &g1 = core::detail::DefaultCPUGenerator(); + const auto &g2 = core::detail::GetDefaultGenerator(CPUDevice()); + EXPECT_EQ(g1.UnsafeGetImpl(), g2.UnsafeGetImpl()); +} + +TEST(GeneratorDefaultTest, GetGeneratorOrDefaultPrefersExplicit) { + auto explicit_gen = MakeCPUGenerator(123); + auto *explicit_impl = core::GetGeneratorOrDefault(explicit_gen, CPUDevice()); + EXPECT_EQ(explicit_impl, explicit_gen.UnsafeGetImpl()); + + auto *default_impl = core::GetGeneratorOrDefault(std::nullopt, CPUDevice()); + EXPECT_EQ(default_impl, core::detail::DefaultCPUGenerator().UnsafeGetImpl()); +} + +TEST(GeneratorDefaultTest, GlobalManualSeedResetsCPUDefault) { + core::ManualSeed(555); + const auto state = core::detail::DefaultCPUGenerator().GetState(); + auto first = core::detail::DefaultCPUGenerator(); + auto first_values = DrawCPU(first, 8); + + core::ManualSeed(555); + auto second = core::detail::DefaultCPUGenerator(); + auto second_values = DrawCPU(second, 8); + + EXPECT_EQ(first_values, second_values); + + auto restored = MakeCPUGenerator(1); + restored.SetState(state); + EXPECT_EQ(restored.CurrentSeed(), 555u); +} + +TEST(GeneratorDefaultTest, DefaultGeneratorStateRoundTrip) { + core::ManualSeed(123); + + auto gen = core::detail::DefaultCPUGenerator(); + + auto state = gen.GetState(); + + auto seq1 = DrawCPU(gen, 32); + + gen.SetState(state); + + auto seq2 = DrawCPU(gen, 32); + + EXPECT_EQ(seq1, seq2); +} + +#ifdef USE_CUDA + +TEST(GeneratorDefaultTest, FutureDefaultGeneratorUsesLastManualSeed) { + core::ManualSeed(2468); + const auto &cuda_gen = core::detail::DefaultCUDAGenerator(7); + EXPECT_EQ(cuda_gen.CurrentSeed(), 2468u); + EXPECT_EQ(cuda_gen.GetDevice(), Device(Device::DeviceType::kCUDA, 7)); +} + +TEST(GeneratorDefaultTest, ManualSeedResetsExistingCUDADefault) { + core::ManualSeed(100); + const auto &cuda5 = core::detail::DefaultCUDAGenerator(5); + EXPECT_EQ(cuda5.CurrentSeed(), 100u); + + core::ManualSeed(200); + EXPECT_EQ(cuda5.CurrentSeed(), 200u); + EXPECT_EQ(cuda5.GetOffset(), 0u); +} + +TEST(GeneratorDefaultTest, MultiDeviceIndependence) { + core::ManualSeed(777); + const auto &cuda3 = core::detail::DefaultCUDAGenerator(3); + const auto &cuda4 = core::detail::DefaultCUDAGenerator(4); + + EXPECT_NE(cuda3.UnsafeGetImpl(), cuda4.UnsafeGetImpl()); + EXPECT_EQ(cuda3.GetDevice(), Device(Device::DeviceType::kCUDA, 3)); + EXPECT_EQ(cuda4.GetDevice(), Device(Device::DeviceType::kCUDA, 4)); + + { + auto *impl3 = cuda3.Get(); + std::lock_guard lock(impl3->mutex_); + impl3->PhiloxEngineInputs(42); + } + EXPECT_EQ(cuda3.GetOffset(), 42u); + EXPECT_EQ(cuda4.GetOffset(), 0u); +} + +TEST(GeneratorDefaultTest, DefaultCUDAGeneratorStateRoundTrip) { + core::ManualSeed(321); + + auto gen = core::detail::DefaultCUDAGenerator(0); + + auto state = gen.GetState(); + + auto *impl = gen.Get(); + + { + std::lock_guard lock(impl->mutex_); + impl->PhiloxEngineInputs(128); + } + + gen.SetState(state); + + EXPECT_EQ(gen.GetOffset(), 0u); +} + +#endif // USE_CUDA diff --git a/tests/generator/test_generator_operators.cc b/tests/generator/test_generator_operators.cc new file mode 100644 index 00000000..55497afd --- /dev/null +++ b/tests/generator/test_generator_operators.cc @@ -0,0 +1,445 @@ +// Generator end-to-end operator integration tests. +// +// These tests verify that random number generation operators (Uniform, Normal, +// KaimingUniform, and Dropout) correctly interact with the Generator mechanism +// on both CPU and CUDA backends, leveraging GTest parameterization. + +#include "infini_train/include/core/cpu_generator.h" +#include "infini_train/include/core/generator.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/datatype.h" +#include "infini_train/include/device.h" +#include "infini_train/include/nn/init.h" +#include "infini_train/include/tensor.h" +#include "tests/common/test_utils.h" +#include +#include +#include + +#ifdef USE_CUDA +#include "infini_train/include/core/cuda_generator.h" +#endif + +using namespace infini_train; + +namespace { +Device CPUDevice() { return Device(Device::DeviceType::kCPU, 0); } +} // namespace + +class GeneratorOperatorsTest : public infini_train::test::InfiniTrainTest { +protected: + std::shared_ptr MakeFloatTensor(int64_t n) { + return std::make_shared(std::vector{n}, DataType::kFLOAT32, GetDevice()); + } + + std::shared_ptr MakeOnesTensor(int64_t n) { + auto tensor = std::make_shared(std::vector{n}, DataType::kFLOAT32, GetDevice()); + nn::init::Ones(tensor); + return tensor; + } + + std::vector ReadFloats(const std::shared_ptr &tensor) { + const int64_t n = tensor->NumElements(); + std::vector host_data(n); + auto impl = core::GetDeviceGuardImpl(GetDevice().type()); + core::DeviceGuard guard(GetDevice()); + impl->MemcpyAsync(host_data.data(), tensor->DataPtr(), n * sizeof(float), + GetDevice().IsCPU() ? core::MemcpyKind::kD2D : core::MemcpyKind::kD2H, + impl->GetStream(GetDevice())); + impl->SynchronizeStream(impl->GetStream(GetDevice())); + return host_data; + } + + core::Generator MakeDeviceGenerator(uint64_t seed) { + if (GetDevice().IsCPU()) { + return core::Generator(std::make_shared(seed)); + } else { +#ifdef USE_CUDA + return core::Generator(std::make_shared(GetDevice().index(), seed)); +#else + LOG(FATAL) << "CUDA generator requested but USE_CUDA is not defined"; +#endif + } + } +}; + +// ============================================================================ +// Initializer Tests (Uniform, Normal, KaimingUniform) +// ============================================================================ + +TEST_P(GeneratorOperatorsTest, UniformReproducibleWithSameSeed) { + auto t1 = MakeFloatTensor(64); + auto t2 = MakeFloatTensor(64); + + auto g1 = MakeDeviceGenerator(2024); + auto g2 = MakeDeviceGenerator(2024); + + nn::init::Uniform(t1, -1.0f, 1.0f, g1); + nn::init::Uniform(t2, -1.0f, 1.0f, g2); + + EXPECT_EQ(ReadFloats(t1), ReadFloats(t2)); +} + +TEST_P(GeneratorOperatorsTest, UniformDiffersWithDifferentSeed) { + auto t1 = MakeFloatTensor(64); + auto t2 = MakeFloatTensor(64); + + auto g1 = MakeDeviceGenerator(1); + auto g2 = MakeDeviceGenerator(2); + + nn::init::Uniform(t1, 0.0f, 1.0f, g1); + nn::init::Uniform(t2, 0.0f, 1.0f, g2); + + EXPECT_NE(ReadFloats(t1), ReadFloats(t2)); +} + +TEST_P(GeneratorOperatorsTest, RepeatedCallsAdvanceGeneratorState) { + auto t1 = MakeFloatTensor(32); + auto t2 = MakeFloatTensor(32); + + auto gen = MakeDeviceGenerator(7); + nn::init::Uniform(t1, 0.0f, 1.0f, gen); + nn::init::Uniform(t2, 0.0f, 1.0f, gen); + + EXPECT_NE(ReadFloats(t1), ReadFloats(t2)); +} + +TEST_P(GeneratorOperatorsTest, NormalReproducibleWithSameSeed) { + auto t1 = MakeFloatTensor(128); + auto t2 = MakeFloatTensor(128); + + auto g1 = MakeDeviceGenerator(99); + auto g2 = MakeDeviceGenerator(99); + + nn::init::Normal(t1, 0.0f, 1.0f, g1); + nn::init::Normal(t2, 0.0f, 1.0f, g2); + + EXPECT_EQ(ReadFloats(t1), ReadFloats(t2)); +} + +TEST_P(GeneratorOperatorsTest, KaimingUniformReproducibleWithSameSeed) { + auto t1 = std::make_shared(std::vector{8, 16}, DataType::kFLOAT32, GetDevice()); + auto t2 = std::make_shared(std::vector{8, 16}, DataType::kFLOAT32, GetDevice()); + + auto g1 = MakeDeviceGenerator(31337); + auto g2 = MakeDeviceGenerator(31337); + + nn::init::KaimingUniform(t1, 0.0f, nn::init::KaimingMode::kFanIn, nn::init::NonLinearityType::kReLU, g1); + nn::init::KaimingUniform(t2, 0.0f, nn::init::KaimingMode::kFanIn, nn::init::NonLinearityType::kReLU, g2); + + EXPECT_EQ(ReadFloats(t1), ReadFloats(t2)); +} + +TEST_P(GeneratorOperatorsTest, TensorUniformReproducibleWithSameSeed) { + auto t1 = MakeFloatTensor(48); + auto t2 = MakeFloatTensor(48); + + auto g1 = MakeDeviceGenerator(555); + auto g2 = MakeDeviceGenerator(555); + + t1->Uniform(-2.0f, 2.0f, g1); + t2->Uniform(-2.0f, 2.0f, g2); + + EXPECT_EQ(ReadFloats(t1), ReadFloats(t2)); +} + +TEST_P(GeneratorOperatorsTest, DefaultGeneratorReproducibleViaManualSeed) { + auto t1 = MakeFloatTensor(64); + auto t2 = MakeFloatTensor(64); + + core::ManualSeed(20240601); + nn::init::Uniform(t1, 0.0f, 1.0f); + + core::ManualSeed(20240601); + nn::init::Uniform(t2, 0.0f, 1.0f); + + EXPECT_EQ(ReadFloats(t1), ReadFloats(t2)); +} + +TEST_P(GeneratorOperatorsTest, ExplicitGeneratorDoesNotConsumeDefault) { + auto baseline = MakeFloatTensor(32); + core::ManualSeed(2024); + nn::init::Uniform(baseline, 0.0f, 1.0f); + + auto with_explicit = MakeFloatTensor(32); + core::ManualSeed(2024); + auto explicit_gen = MakeDeviceGenerator(123456); + auto scratch = MakeFloatTensor(32); + nn::init::Uniform(scratch, 0.0f, 1.0f, explicit_gen); + nn::init::Uniform(with_explicit, 0.0f, 1.0f); + + EXPECT_EQ(ReadFloats(baseline), ReadFloats(with_explicit)); +} + +TEST_P(GeneratorOperatorsTest, KaimingUniformDefaultGeneratorReproducible) { + auto t1 = std::make_shared(std::vector{16, 32}, DataType::kFLOAT32, GetDevice()); + auto t2 = std::make_shared(std::vector{16, 32}, DataType::kFLOAT32, GetDevice()); + + core::ManualSeed(9999); + nn::init::KaimingUniform(t1); + + core::ManualSeed(9999); + nn::init::KaimingUniform(t2); + + EXPECT_EQ(ReadFloats(t1), ReadFloats(t2)); +} + +TEST_P(GeneratorOperatorsTest, NormalDefaultGeneratorReproducible) { + auto t1 = MakeFloatTensor(128); + auto t2 = MakeFloatTensor(128); + + core::ManualSeed(7777); + nn::init::Normal(t1, 0.0f, 0.02f); + + core::ManualSeed(7777); + nn::init::Normal(t2, 0.0f, 0.02f); + + EXPECT_EQ(ReadFloats(t1), ReadFloats(t2)); +} + +TEST_P(GeneratorOperatorsTest, RejectsNonFloat32Tensor) { + auto int_tensor = std::make_shared(std::vector{8}, DataType::kINT32, GetDevice()); + auto gen = MakeDeviceGenerator(1); + EXPECT_DEATH(nn::init::Uniform(int_tensor, 0.0f, 1.0f, gen), "only supports FLOAT32"); +} + +TEST_P(GeneratorOperatorsTest, RejectsGeneratorDeviceMismatch) { +#ifdef USE_CUDA + if (GetDevice().IsCPU()) { + auto cpu_tensor = std::make_shared(std::vector{8}, DataType::kFLOAT32, CPUDevice()); + core::Generator cuda_gen(std::make_shared(0, 1)); + EXPECT_DEATH(nn::init::Uniform(cpu_tensor, 0.0f, 1.0f, cuda_gen), "device type mismatch"); + } else { + auto cuda_tensor = std::make_shared(std::vector{8}, DataType::kFLOAT32, GetDevice()); + core::Generator cpu_gen(std::make_shared(1)); + EXPECT_DEATH(nn::init::Uniform(cuda_tensor, 0.0f, 1.0f, cpu_gen), "device type mismatch"); + } +#endif +} + +// ============================================================================ +// Dropout Tests +// ============================================================================ + +TEST_P(GeneratorOperatorsTest, DropoutReproducibleWithSameSeed) { + auto t1 = MakeOnesTensor(256); + auto t2 = MakeOnesTensor(256); + + auto g1 = MakeDeviceGenerator(2024); + auto g2 = MakeDeviceGenerator(2024); + + nn::init::Dropout(t1, 0.5f, g1); + nn::init::Dropout(t2, 0.5f, g2); + + EXPECT_EQ(ReadFloats(t1), ReadFloats(t2)); +} + +TEST_P(GeneratorOperatorsTest, DropoutZerosAndScalesSurvivors) { + const float p = 0.3f; + const float scale = 1.0f / (1.0f - p); + + auto tensor = MakeOnesTensor(1024); + auto gen = MakeDeviceGenerator(7); + nn::init::Dropout(tensor, p, gen); + + int zeros = 0; + int survivors = 0; + for (float v : ReadFloats(tensor)) { + if (v == 0.0f) { + ++zeros; + } else { + EXPECT_FLOAT_EQ(v, scale); + ++survivors; + } + } + EXPECT_GT(zeros, 0); + EXPECT_GT(survivors, 0); +} + +TEST_P(GeneratorOperatorsTest, DropoutDifferentSeedProducesDifferentMask) { + auto t1 = MakeOnesTensor(256); + auto t2 = MakeOnesTensor(256); + + auto g1 = MakeDeviceGenerator(1); + auto g2 = MakeDeviceGenerator(2); + + nn::init::Dropout(t1, 0.5f, g1); + nn::init::Dropout(t2, 0.5f, g2); + + EXPECT_NE(ReadFloats(t1), ReadFloats(t2)); +} + +TEST_P(GeneratorOperatorsTest, DropoutRepeatedCallsAdvanceGeneratorState) { + auto t1 = MakeOnesTensor(256); + auto t2 = MakeOnesTensor(256); + + auto gen = MakeDeviceGenerator(11); + nn::init::Dropout(t1, 0.5f, gen); + nn::init::Dropout(t2, 0.5f, gen); + + EXPECT_NE(ReadFloats(t1), ReadFloats(t2)); +} + +TEST_P(GeneratorOperatorsTest, DropoutUsesDefaultGeneratorReproducibly) { + auto t1 = MakeOnesTensor(256); + auto t2 = MakeOnesTensor(256); + + core::ManualSeed(20240601); + nn::init::Dropout(t1, 0.5f); + + core::ManualSeed(20240601); + nn::init::Dropout(t2, 0.5f); + + EXPECT_EQ(ReadFloats(t1), ReadFloats(t2)); +} + +TEST_P(GeneratorOperatorsTest, DropoutBoundaryProbabilityZeroKeepsAllElements) { + auto tensor = MakeOnesTensor(512); + auto gen = MakeDeviceGenerator(3); + nn::init::Dropout(tensor, 0.0f, gen); + + for (float v : ReadFloats(tensor)) { EXPECT_FLOAT_EQ(v, 1.0f); } +} + +TEST_P(GeneratorOperatorsTest, DropoutBoundaryProbabilityNearOneDropsMost) { + const float p = 0.99f; + auto tensor = MakeOnesTensor(4096); + auto gen = MakeDeviceGenerator(123); + nn::init::Dropout(tensor, p, gen); + + int zeros = 0; + for (float v : ReadFloats(tensor)) { + if (v == 0.0f) { + ++zeros; + } + } + EXPECT_GT(zeros, 4096 * 9 / 10); +} + +TEST_P(GeneratorOperatorsTest, DropoutEmptyTensorIsNoOp) { + auto tensor = MakeFloatTensor(0); + auto gen = MakeDeviceGenerator(321); + + nn::init::Dropout(tensor, 0.5f, gen); + + EXPECT_EQ(tensor->NumElements(), 0); + EXPECT_EQ(tensor->Dims(), (std::vector{0})); +} + +TEST_P(GeneratorOperatorsTest, DropoutRejectsIllegalProbability) { + auto gen = MakeDeviceGenerator(1); + { + auto tensor = MakeOnesTensor(8); + EXPECT_DEATH(nn::init::Dropout(tensor, 1.0f, gen), "probability must be in"); + } + { + auto tensor = MakeOnesTensor(8); + EXPECT_DEATH(nn::init::Dropout(tensor, -0.1f, gen), "probability must be in"); + } +} + +// ============================================================================ + +TEST_P(GeneratorOperatorsTest, DifferentTensorShapesAdvanceOffsetCorrectly) { + SKIP_CPU(); + auto gen = MakeDeviceGenerator(123); + + auto t1 = MakeFloatTensor(13); + auto t2 = MakeFloatTensor(257); + + nn::init::Uniform(t1, 0.f, 1.f, gen); + auto offset1 = gen.GetOffset(); + + nn::init::Uniform(t2, 0.f, 1.f, gen); + auto offset2 = gen.GetOffset(); + + EXPECT_GT(offset1, 0u); + EXPECT_GT(offset2, offset1); +} + +TEST_P(GeneratorOperatorsTest, RestoreGeneratorStateRestoresSequence) { + auto gen = MakeDeviceGenerator(777); + + auto state = gen.GetState(); + + auto t1 = MakeFloatTensor(64); + auto t2 = MakeFloatTensor(64); + + nn::init::Uniform(t1, 0, 1, gen); + + gen.SetState(state); + + nn::init::Uniform(t2, 0, 1, gen); + + EXPECT_EQ(ReadFloats(t1), ReadFloats(t2)); +} + +// ============================================================================ +// Generator Offset Tests +// ============================================================================ +TEST_P(GeneratorOperatorsTest, UniformOffsetMatchesExpected) { + SKIP_CPU(); + auto gen = MakeDeviceGenerator(123); + + auto tensor = MakeFloatTensor(1000); + + nn::init::Uniform(tensor, 0.f, 1.f, gen); + + EXPECT_GT(gen.GetOffset(), 0u); +} +TEST_P(GeneratorOperatorsTest, NormalOffsetMatchesExpected) { + SKIP_CPU(); + auto gen = MakeDeviceGenerator(123); + + auto tensor = MakeFloatTensor(1000); + + nn::init::Normal(tensor, 0.f, 1.f, gen); + + EXPECT_GT(gen.GetOffset(), 0u); +} +// ============================================================================ + +TEST_P(GeneratorOperatorsTest, DropoutAdvancesGeneratorOffset) { + SKIP_CPU(); + auto gen = MakeDeviceGenerator(123); + + auto tensor = MakeOnesTensor(1000); + + EXPECT_EQ(gen.GetOffset(), 0u); + + nn::init::Dropout(tensor, 0.5f, gen); + + EXPECT_GT(gen.GetOffset(), 0u); +} +TEST_P(GeneratorOperatorsTest, UniformUsesSameGeneratorAcrossTensorSizes) { + SKIP_CPU(); + auto gen = MakeDeviceGenerator(123); + + auto a = MakeFloatTensor(17); + auto b = MakeFloatTensor(31); + + nn::init::Uniform(a, 0.f, 1.f, gen); + + auto offset_after_first = gen.GetOffset(); + + nn::init::Uniform(b, 0.f, 1.f, gen); + + EXPECT_GT(gen.GetOffset(), offset_after_first); +} +TEST_P(GeneratorOperatorsTest, DropoutUsesSameGeneratorAcrossTensorSizes) { + SKIP_CPU(); + auto gen = MakeDeviceGenerator(123); + + auto a = MakeOnesTensor(17); + auto b = MakeOnesTensor(31); + + nn::init::Dropout(a, 0.5f, gen); + + auto offset_after_first = gen.GetOffset(); + + nn::init::Dropout(b, 0.5f, gen); + + EXPECT_GT(gen.GetOffset(), offset_after_first); +} + +INFINI_TRAIN_REGISTER_TEST(GeneratorOperatorsTest); diff --git a/tests/generator/test_generator_thread_safety.cc b/tests/generator/test_generator_thread_safety.cc new file mode 100644 index 00000000..0b29c050 --- /dev/null +++ b/tests/generator/test_generator_thread_safety.cc @@ -0,0 +1,200 @@ +#include "infini_train/include/core/cpu_generator.h" +#include "infini_train/include/core/generator.h" +#include "infini_train/include/device.h" +#include "gtest/gtest.h" +#include +#include +#include + +using namespace infini_train; + +namespace { + +core::Generator MakeCPUGenerator(uint64_t seed) { + return core::Generator(std::make_shared(seed)); +} + +std::vector DrawCPU(core::Generator &gen, int64_t n) { + auto *impl = gen.Get(); + std::lock_guard lock(gen.Mutex()); + std::vector values(static_cast(n)); + auto &engine = impl->Engine(); + for (auto &v : values) { v = engine(); } + return values; +} + +} // namespace + +TEST(GeneratorThreadSafety, MultipleThreadsDraw) { + auto gen = MakeCPUGenerator(42); + constexpr int kNumThreads = 4; + constexpr int kDrawsPerThread = 1000; + + auto worker = [&gen]() { + for (int i = 0; i < kDrawsPerThread; ++i) { + auto *impl = gen.Get(); + std::lock_guard lock(gen.Mutex()); + impl->Engine()(); + } + }; + + std::vector threads; + for (int i = 0; i < kNumThreads; ++i) { threads.emplace_back(worker); } + for (auto &t : threads) { t.join(); } + + EXPECT_TRUE(gen.Defined()); + + auto state = gen.GetState(); + auto values1 = DrawCPU(gen, 32); + gen.SetState(state); + auto values2 = DrawCPU(gen, 32); + EXPECT_EQ(values1, values2); +} + +TEST(GeneratorThreadSafety, StateConcurrent) { + auto gen = MakeCPUGenerator(123); + constexpr int kNumThreads = 4; + + auto worker = [&gen]() { + for (int i = 0; i < 50; ++i) { + auto state = gen.GetState(); + gen.SetState(state); + + auto *impl = gen.Get(); + std::lock_guard lock(gen.Mutex()); + impl->Engine()(); + } + }; + + std::vector threads; + for (int i = 0; i < kNumThreads; ++i) { threads.emplace_back(worker); } + for (auto &t : threads) { t.join(); } + EXPECT_TRUE(gen.Defined()); + + auto state = gen.GetState(); + auto values1 = DrawCPU(gen, 32); + gen.SetState(state); + auto values2 = DrawCPU(gen, 32); + EXPECT_EQ(values1, values2); +} + +TEST(GeneratorThreadSafety, ManualSeedConcurrent) { + auto gen = MakeCPUGenerator(777); + constexpr int kNumThreads = 4; + + auto worker = [&gen](uint64_t thread_seed) { + for (int i = 0; i < 50; ++i) { + gen.ManualSeed(thread_seed + static_cast(i)); + auto s = gen.CurrentSeed(); + (void)s; + } + }; + + std::vector threads; + for (int i = 0; i < kNumThreads; ++i) { threads.emplace_back(worker, static_cast(i) * 1000); } + for (auto &t : threads) { t.join(); } + EXPECT_TRUE(gen.Defined()); + + gen.ManualSeed(42); + auto values1 = DrawCPU(gen, 32); + + auto gen2 = MakeCPUGenerator(42); + auto values2 = DrawCPU(gen2, 32); + EXPECT_EQ(values1, values2); +} + +TEST(GeneratorThreadSafety, DefaultGeneratorConcurrent) { + constexpr int kNumThreads = 4; + + auto worker = [](uint64_t thread_seed) { + for (int i = 0; i < 50; ++i) { + core::ManualSeed(thread_seed + static_cast(i)); + auto gen = core::detail::DefaultCPUGenerator(); + auto state = gen.GetState(); + + auto *impl = gen.Get(); + std::lock_guard lock(gen.Mutex()); + impl->Engine()(); + } + }; + + std::vector threads; + for (int i = 0; i < kNumThreads; ++i) { threads.emplace_back(worker, static_cast(i) * 1000); } + for (auto &t : threads) { t.join(); } + EXPECT_TRUE(core::detail::DefaultCPUGenerator().Defined()); + + core::ManualSeed(99); + auto gen = core::detail::DefaultCPUGenerator(); + auto state = gen.GetState(); + auto values1 = DrawCPU(gen, 32); + gen.SetState(state); + auto values2 = DrawCPU(gen, 32); + EXPECT_EQ(values1, values2); +} + +TEST(GeneratorThreadSafety, ConcurrentSeedAndDrawConsistency) { + constexpr int kNumThreads = 8; + + auto gen = MakeCPUGenerator(0); + + auto worker = [&gen](int id) { + if (id % 2 == 0) { + for (int i = 0; i < 100; ++i) { + gen.ManualSeed(static_cast(id) * 1000 + static_cast(i)); + } + } else { + for (int i = 0; i < 100; ++i) { + auto *impl = gen.Get(); + std::lock_guard lock(gen.Mutex()); + (void)impl->Engine()(); + } + } + }; + + std::vector threads; + for (int i = 0; i < kNumThreads; ++i) { threads.emplace_back(worker, i); } + for (auto &t : threads) { t.join(); } + EXPECT_TRUE(gen.Defined()); + + gen.ManualSeed(12345); + auto state = gen.GetState(); + auto values1 = DrawCPU(gen, 64); + gen.SetState(state); + auto values2 = DrawCPU(gen, 64); + EXPECT_EQ(values1, values2); +} + +TEST(GeneratorThreadSafety, CloneConcurrentDraw) { + auto gen = MakeCPUGenerator(123); + + auto cloned = gen.Clone(); + + std::thread t1([&]() { DrawCPU(gen, 1000); }); + + std::thread t2([&]() { DrawCPU(cloned, 1000); }); + + t1.join(); + + t2.join(); + + EXPECT_TRUE(gen.Defined()); + + EXPECT_TRUE(cloned.Defined()); +} + +TEST(GeneratorThreadSafety, ConcurrentClone) { + auto gen = MakeCPUGenerator(42); + + constexpr int kThreads = 8; + + std::vector threads; + + for (int i = 0; i < kThreads; ++i) { + threads.emplace_back([&]() { + auto cloned = gen.Clone(); + EXPECT_TRUE(cloned.Defined()); + }); + } + + for (auto &t : threads) { t.join(); } +} \ No newline at end of file