From 921e826fb66d195e3acb5286ed33cfc66099ece0 Mon Sep 17 00:00:00 2001 From: Kinorw Date: Wed, 8 Jul 2026 22:26:13 +0800 Subject: [PATCH 1/7] feat(generator): add Generator and GeneratorImpl classes with basic API tests --- infini_train/include/core/generator.h | 178 ++++++++++++++++++++++++++ tests/CMakeLists.txt | 3 + tests/generator/CMakeLists.txt | 23 ++++ tests/generator/test_generator_api.cc | 150 ++++++++++++++++++++++ 4 files changed, 354 insertions(+) create mode 100644 infini_train/include/core/generator.h create mode 100644 tests/generator/CMakeLists.txt create mode 100644 tests/generator/test_generator_api.cc diff --git a/infini_train/include/core/generator.h b/infini_train/include/core/generator.h new file mode 100644 index 00000000..d85b4375 --- /dev/null +++ b/infini_train/include/core/generator.h @@ -0,0 +1,178 @@ +#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/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..fb4f09d0 --- /dev/null +++ b/tests/generator/CMakeLists.txt @@ -0,0 +1,23 @@ +# ========================================================================== +# Generator (random number generator) tests +# ========================================================================== + +file(GLOB GENERATOR_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/test_generator_*.cc) + +# CPU tests: run all tests that do NOT have "CUDA" in their test suite name +# or test parameterization prefix. +infini_train_add_test(test_generator + SOURCES ${GENERATOR_SOURCES} + LABELS cpu + TEST_FILTER "-*CUDA*" +) + +# CUDA tests: run tests that DO have "CUDA" in their test suite name +# or test parameterization prefix. +if(USE_CUDA) + infini_train_add_test(test_generator_cuda + SOURCES ${GENERATOR_SOURCES} + LABELS cuda + TEST_FILTER "*CUDA*" + ) +endif() diff --git a/tests/generator/test_generator_api.cc b/tests/generator/test_generator_api.cc new file mode 100644 index 00000000..006aeade --- /dev/null +++ b/tests/generator/test_generator_api.cc @@ -0,0 +1,150 @@ +// 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: + 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 new_seed = gen.Seed(); + EXPECT_EQ(new_seed, 201u); + EXPECT_EQ(gen.CurrentSeed(), 201u); + + 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)); +} From 69993b1f9f30e57574d4050b9b04ddb14dbf1454 Mon Sep 17 00:00:00 2001 From: Kinorw Date: Wed, 8 Jul 2026 22:26:34 +0800 Subject: [PATCH 2/7] feat(generator): implement CPUGenerator with state management and unit tests --- infini_train/include/core/cpu_generator.h | 57 ++++++++++ infini_train/src/core/cpu_generator.cc | 94 ++++++++++++++++ tests/generator/test_generator_cpu.cc | 126 ++++++++++++++++++++++ 3 files changed, 277 insertions(+) create mode 100644 infini_train/include/core/cpu_generator.h create mode 100644 infini_train/src/core/cpu_generator.cc create mode 100644 tests/generator/test_generator_cpu.cc diff --git a/infini_train/include/core/cpu_generator.h b/infini_train/include/core/cpu_generator.h new file mode 100644 index 00000000..197d0506 --- /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/src/core/cpu_generator.cc b/infini_train/src/core/cpu_generator.cc new file mode 100644 index 00000000..c5e70a2c --- /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/tests/generator/test_generator_cpu.cc b/tests/generator/test_generator_cpu.cc new file mode 100644 index 00000000..2198a40b --- /dev/null +++ b/tests/generator/test_generator_cpu.cc @@ -0,0 +1,126 @@ +// CPU Generator unit tests. + +#include "gtest/gtest.h" +#include "infini_train/include/core/cpu_generator.h" +#include "infini_train/include/core/generator.h" +#include "infini_train/include/device.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; +} + +} // 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, RejectsMalformedCPUState) { + auto gen = MakeCPUGenerator(1); + EXPECT_DEATH(gen.SetState(std::vector{1, 2, 3}), "Invalid CPU generator state"); +} + +TEST(CPUGeneratorTest, RejectsForeignBackendState) { + auto gen = MakeCPUGenerator(1); + std::vector foreign(32, 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"); +} From 984ced00b1954ff1a99c3d301a7ae97741d63950 Mon Sep 17 00:00:00 2001 From: Kinorw Date: Wed, 8 Jul 2026 22:38:47 +0800 Subject: [PATCH 3/7] feat(generator): add CUDAGenerator implementation with state management and unit tests --- infini_train/include/core/cuda_generator.h | 63 +++++++++++++ infini_train/src/core/cuda_generator.cc | 100 +++++++++++++++++++++ tests/generator/test_generator_cuda.cc | 98 ++++++++++++++++++++ 3 files changed, 261 insertions(+) create mode 100644 infini_train/include/core/cuda_generator.h create mode 100644 infini_train/src/core/cuda_generator.cc create mode 100644 tests/generator/test_generator_cuda.cc 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/src/core/cuda_generator.cc b/infini_train/src/core/cuda_generator.cc new file mode 100644 index 00000000..a5634363 --- /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/tests/generator/test_generator_cuda.cc b/tests/generator/test_generator_cuda.cc new file mode 100644 index 00000000..1ee34322 --- /dev/null +++ b/tests/generator/test_generator_cuda.cc @@ -0,0 +1,98 @@ +// CUDA Generator unit tests. + +#include "gtest/gtest.h" +#include "infini_train/include/core/generator.h" +#include "infini_train/include/device.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, 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); +} + +#endif // USE_CUDA From fe7cbc552bf289eaa3ae41ba1b43731c56cee7d9 Mon Sep 17 00:00:00 2001 From: Kinorw Date: Wed, 8 Jul 2026 22:39:03 +0800 Subject: [PATCH 4/7] feat(generator): implement default generator management and ManualSeed tests --- infini_train/src/core/generator.cc | 105 +++++++++++++++++++++ tests/generator/test_generator_default.cc | 110 ++++++++++++++++++++++ 2 files changed, 215 insertions(+) create mode 100644 infini_train/src/core/generator.cc create mode 100644 tests/generator/test_generator_default.cc diff --git a/infini_train/src/core/generator.cc b/infini_train/src/core/generator.cc new file mode 100644 index 00000000..58fbdd35 --- /dev/null +++ b/infini_train/src/core/generator.cc @@ -0,0 +1,105 @@ +#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/tests/generator/test_generator_default.cc b/tests/generator/test_generator_default.cc new file mode 100644 index 00000000..9948a689 --- /dev/null +++ b/tests/generator/test_generator_default.cc @@ -0,0 +1,110 @@ +// Default generator management and ManualSeed tests. + +#include "gtest/gtest.h" +#include "infini_train/include/core/generator.h" +#include "infini_train/include/device.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); +} + +#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); +} + +#endif // USE_CUDA From cc900218b81b31289ed0421ea95654c11f0cae20 Mon Sep 17 00:00:00 2001 From: Kinorw Date: Fri, 10 Jul 2026 01:57:58 +0800 Subject: [PATCH 5/7] feat(generator): implement distribution infrastructure and integrate with tensor initialization --- infini_train/include/core/distribution.h | 138 ++++++++ infini_train/include/nn/init.h | 11 +- infini_train/include/tensor.h | 3 +- infini_train/src/core/distribution.cc | 20 ++ infini_train/src/nn/init.cc | 103 +++--- infini_train/src/tensor.cc | 2 +- tests/generator/test_generator_operators.cc | 335 ++++++++++++++++++++ 7 files changed, 553 insertions(+), 59 deletions(-) create mode 100644 infini_train/include/core/distribution.h create mode 100644 infini_train/src/core/distribution.cc create mode 100644 tests/generator/test_generator_operators.cc diff --git a/infini_train/include/core/distribution.h b/infini_train/include/core/distribution.h new file mode 100644 index 00000000..9b140b8d --- /dev/null +++ b/infini_train/include/core/distribution.h @@ -0,0 +1,138 @@ +#pragma once + +// distribution.h — Minimal distribution template infrastructure. +// +// Mirrors PyTorch's DistributionTemplates.h / cpu_DistributionTemplates.h / +// cuda_DistributionTemplates.h three-layer design, simplified for InfiniTrain's +// current needs: +// +// 1. distribution_impl — parameter validation + backend dispatch +// 2. cpu::*_kernel — CPU backend: locks generator, fills elements +// 3. cuda::*_kernel — CUDA backend: host-side fallback for now +// +// The key improvement over the old FillBuffer is: +// - Distribution logic is decoupled from init.cc consumers +// - Generator resolution (explicit vs default) follows PyTorch's +// check_generator / GetGeneratorOrDefault pattern +// - Each backend owns its locking strategy +// - Ready for future device-kernel extension on CUDA + +#include +#include +#include +#include +#include + +#include "infini_train/include/core/cpu_generator.h" +#include "infini_train/include/core/generator.h" +#ifdef USE_CUDA +#include "infini_train/include/core/cuda_generator.h" +#endif +#include "infini_train/include/device.h" + +namespace infini_train::core::distribution { + +// ============================== CPU backend ============================== + +namespace cpu { + +// Fills `buffer` with values drawn from `dist`, using the CPU generator's +// engine directly. The generator mutex is held for the entire draw loop +// because Engine() returns a reference to the generator's internal state. +// +// Corresponds to PyTorch's cpu::uniform_kernel / cpu::normal_kernel pattern, +// simplified to operate on a flat float buffer instead of a TensorIterator. +template +void FillKernel(float *data, int64_t numel, 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 cpu + +// ============================== CUDA backend ============================= + +#ifdef USE_CUDA +namespace cuda { + +// Host-side fallback: generates random values on the host using a temporary +// mt19937_64 engine derived from the CUDA generator's (seed, offset) pair, +// then the caller copies the result to the device. +// +// A future device-kernel path will generate values directly on the GPU via +// Philox (mirroring PyTorch's distribution_nullary_kernel). +// +// Locking strategy matches PyTorch: only the offset reservation (PhiloxEngine- +// Inputs) touches impl state; the subsequent host-side generation uses purely +// local variables and runs outside the impl lock to reduce contention. +template +void FillKernel(float *data, int64_t numel, CUDAGeneratorImpl *impl, Dist dist) { + uint64_t seed, offset; + { + std::lock_guard lock(impl->mutex_); + std::tie(seed, offset) = impl->PhiloxEngineInputs(static_cast(numel)); + } + // Mix a backend-specific constant so that the derived engine produces a + // sequence distinct from the CPU generator even when seed and offset + // values coincide (e.g. both backends seeded identically via ManualSeed). + constexpr uint64_t kCUDABackendTag = 0xD2511F53D9E08B47ULL; + std::mt19937_64 engine(seed ^ ((offset + kCUDABackendTag) * 0x9E3779B97F4A7C15ULL)); + for (int64_t i = 0; i < numel; ++i) { + data[i] = dist(engine); + } +} + +} // namespace cuda +#endif // USE_CUDA + +// ======================== Distribution dispatch ========================== +// +// These mirror PyTorch's DistributionTemplates.h top-level functions +// (normal_impl_, uniform_impl_, etc.) which validate parameters, resolve the +// generator, and dispatch to the correct backend kernel. + +// Fills `buffer` by drawing from the given distribution, using the resolved +// generator for the target device. Replaces the old FillBuffer in init.cc. +// +// Flow: +// 1. Resolve generator (explicit or device default) +// 2. Dispatch to cpu::FillKernel or cuda::FillKernel +// 3. Backend kernel handles locking and element-wise generation +template +void FillBuffer(std::vector &buffer, const Device &device, + const std::optional &generator, Dist dist) { + auto *data = buffer.data(); + const auto numel = static_cast(buffer.size()); + + if (device.IsCPU()) { + auto *impl = GetGeneratorOrDefault(generator, device); + cpu::FillKernel(data, numel, impl, dist); + return; + } + +#ifdef USE_CUDA + auto *impl = GetGeneratorOrDefault(generator, device); + cuda::FillKernel(data, numel, impl, dist); +#else + LOG(FATAL) << "Random operator on non-CPU device requires CUDA support"; +#endif +} + +// Convenience wrappers matching PyTorch's named distribution templates. + +// uniform_impl_: fills buffer with values from [from, to). +inline void UniformFill(std::vector &buffer, const Device &device, + const std::optional &generator, float from, float to) { + FillBuffer(buffer, device, generator, std::uniform_real_distribution(from, to)); +} + +// normal_impl_: fills buffer with values from N(mean, std). +inline void NormalFill(std::vector &buffer, const Device &device, + const std::optional &generator, float mean, float std) { + FillBuffer(buffer, device, generator, std::normal_distribution(mean, std)); +} + +} // namespace infini_train::core::distribution 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/distribution.cc b/infini_train/src/core/distribution.cc new file mode 100644 index 00000000..5c0681b8 --- /dev/null +++ b/infini_train/src/core/distribution.cc @@ -0,0 +1,20 @@ +// distribution.cc — Non-template distribution utilities. +// +// The core template machinery lives in distribution.h (all inline/template). +// This translation unit exists for: +// 1. Future non-template distribution helpers (e.g. bounds checking) +// 2. Ensuring the header compiles cleanly as a standalone TU +// 3. Providing a home for distribution-related constants if needed + +#include "infini_train/include/core/distribution.h" + +namespace infini_train::core::distribution { + +// Currently all distribution logic is header-only (template + inline). +// This TU is reserved for future non-template distribution utilities, +// such as: +// - check_uniform_bounds (parameter validation, cf. PyTorch's +// CHECK_OUT_OF_BOUNDS / check_from_to_in_range) +// - Distribution registration for multi-dtype dispatch + +} // namespace infini_train::core::distribution diff --git a/infini_train/src/nn/init.cc b/infini_train/src/nn/init.cc index 79b4b48b..025f4308 100644 --- a/infini_train/src/nn/init.cc +++ b/infini_train/src/nn/init.cc @@ -7,55 +7,27 @@ #include #include -#ifdef USE_OMP -#include -#endif - #include "glog/logging.h" +#include "infini_train/include/core/distribution.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/tensor.h" 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) { + 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"; 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 - auto device = tensor->GetDevice(); + core::distribution::FillBuffer(buffer, device, generator, std::normal_distribution(mean, std)); + core::DeviceGuard guard(device); auto impl = core::GetDeviceGuardImpl(device.type()); @@ -113,7 +85,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,26 +100,15 @@ 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::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"; 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 - auto device = tensor->GetDevice(); + core::distribution::FillBuffer(buffer, device, generator, std::uniform_real_distribution(a, b)); core::DeviceGuard guard(device); auto impl = core::GetDeviceGuardImpl(device.type()); @@ -159,6 +120,42 @@ std::shared_ptr Uniform(const std::shared_ptr &tensor, float a, return tensor; } +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"; + + const int64_t num_elements = tensor->NumElements(); + auto device = tensor->GetDevice(); + + // Draw uniform values, then build an inverted dropout mask on the host. + std::vector values(num_elements); + core::distribution::FillBuffer(values, device, generator, std::uniform_real_distribution(0.0f, 1.0f)); + + const float scale = 1.0f / (1.0f - p); + + core::DeviceGuard guard(device); + auto impl = core::GetDeviceGuardImpl(device.type()); + + // Pull current tensor data to host, apply the mask, and write it back. This + // keeps the operator backend-agnostic while still routing randomness through + // the Generator mechanism. + std::vector data(num_elements); + impl->MemcpyAsync(data.data(), tensor->DataPtr(), num_elements * sizeof(float), + device.IsCPU() ? core::MemcpyKind::kD2D : core::MemcpyKind::kD2H, impl->GetStream(device)); + impl->SynchronizeStream(impl->GetStream(device)); + + for (int64_t i = 0; i < num_elements; ++i) { + data[i] = (values[i] < p) ? 0.0f : data[i] * scale; + } + + impl->MemcpyAsync(tensor->DataPtr(), data.data(), num_elements * sizeof(float), + device.IsCPU() ? core::MemcpyKind::kD2D : core::MemcpyKind::kH2D, impl->GetStream(device)); + return tensor; +} + std::shared_ptr Ones(const std::shared_ptr &tensor) { // TODO(dcj): Support other data types later. CHECK_EQ(static_cast(tensor->Dtype()), static_cast(DataType::kFLOAT32)); 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/generator/test_generator_operators.cc b/tests/generator/test_generator_operators.cc new file mode 100644 index 00000000..83285192 --- /dev/null +++ b/tests/generator/test_generator_operators.cc @@ -0,0 +1,335 @@ +// 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 "tests/common/test_utils.h" +#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 +#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, 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"); + } +} + +INFINI_TRAIN_REGISTER_TEST(GeneratorOperatorsTest); From ff606cea7a96432ab7ef02cc8249a7f8053ef6f2 Mon Sep 17 00:00:00 2001 From: Kinorw Date: Sun, 12 Jul 2026 03:41:57 +0800 Subject: [PATCH 6/7] feat(tests): add thread safety tests for CPU generator --- tests/generator/test_generator_api.cc | 42 +++- tests/generator/test_generator_cpu.cc | 63 +++++- .../generator/test_generator_thread_safety.cc | 187 ++++++++++++++++++ 3 files changed, 287 insertions(+), 5 deletions(-) create mode 100644 tests/generator/test_generator_thread_safety.cc diff --git a/tests/generator/test_generator_api.cc b/tests/generator/test_generator_api.cc index 006aeade..b70d0fd4 100644 --- a/tests/generator/test_generator_api.cc +++ b/tests/generator/test_generator_api.cc @@ -20,6 +20,8 @@ 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 { @@ -135,9 +137,10 @@ TEST(GeneratorAPITest, BasicAPIForwarding) { gen.ManualSeed(200); EXPECT_EQ(gen.CurrentSeed(), 200u); + uint64_t old_seed = gen.CurrentSeed(); uint64_t new_seed = gen.Seed(); - EXPECT_EQ(new_seed, 201u); - EXPECT_EQ(gen.CurrentSeed(), 201u); + EXPECT_NE(new_seed, old_seed); + EXPECT_EQ(new_seed, gen.CurrentSeed()); EXPECT_EQ(gen.GetOffset(), 0u); gen.SetOffset(99); @@ -148,3 +151,38 @@ TEST(GeneratorAPITest, BasicAPIForwarding) { 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()); +} diff --git a/tests/generator/test_generator_cpu.cc b/tests/generator/test_generator_cpu.cc index 2198a40b..0baba4c1 100644 --- a/tests/generator/test_generator_cpu.cc +++ b/tests/generator/test_generator_cpu.cc @@ -28,6 +28,19 @@ std::vector DrawCPU(core::Generator &generator, int64_t count) { 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) { @@ -93,14 +106,16 @@ TEST(CPUGeneratorTest, StateRoundTripPreservesSeed) { EXPECT_EQ(other.CurrentSeed(), 0xABCDEFu); } -TEST(CPUGeneratorTest, RejectsMalformedCPUState) { +TEST(CPUGeneratorTest, RejectsStateTooShort) { auto gen = MakeCPUGenerator(1); - EXPECT_DEATH(gen.SetState(std::vector{1, 2, 3}), "Invalid CPU generator state"); + 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(32, 0); + std::vector foreign(gen.GetState().size(), 0); EXPECT_DEATH(gen.SetState(foreign), "backend magic mismatch"); } @@ -124,3 +139,45 @@ TEST(CPUGeneratorTest, CPUOnlySupportsZeroOffset) { 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_thread_safety.cc b/tests/generator/test_generator_thread_safety.cc new file mode 100644 index 00000000..380f9a67 --- /dev/null +++ b/tests/generator/test_generator_thread_safety.cc @@ -0,0 +1,187 @@ +#include "gtest/gtest.h" +#include "infini_train/include/core/generator.h" +#include "infini_train/include/core/cpu_generator.h" +#include "infini_train/include/device.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); +} From 901bf0ec14f3091b334e5d2419e9a785d3408d5c Mon Sep 17 00:00:00 2001 From: Kinorw Date: Sun, 12 Jul 2026 03:44:13 +0800 Subject: [PATCH 7/7] style: format code for consistency and readability across generator files --- infini_train/include/core/cpu_generator.h | 2 +- infini_train/include/core/distribution.h | 34 +++++------ infini_train/include/core/generator.h | 23 +++----- infini_train/src/core/cpu_generator.cc | 6 +- infini_train/src/core/cuda_generator.cc | 6 +- infini_train/src/core/distribution.cc | 2 +- infini_train/src/core/generator.cc | 8 +-- infini_train/src/nn/init.cc | 4 +- tests/generator/test_generator_api.cc | 38 ++++--------- tests/generator/test_generator_cpu.cc | 17 ++---- tests/generator/test_generator_cuda.cc | 6 +- tests/generator/test_generator_default.cc | 10 +--- tests/generator/test_generator_operators.cc | 10 +--- .../generator/test_generator_thread_safety.cc | 56 ++++++------------- 14 files changed, 76 insertions(+), 146 deletions(-) diff --git a/infini_train/include/core/cpu_generator.h b/infini_train/include/core/cpu_generator.h index 197d0506..d34d6329 100644 --- a/infini_train/include/core/cpu_generator.h +++ b/infini_train/include/core/cpu_generator.h @@ -17,7 +17,7 @@ namespace infini_train::core { * 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; diff --git a/infini_train/include/core/distribution.h b/infini_train/include/core/distribution.h index 9b140b8d..03ed2bea 100644 --- a/infini_train/include/core/distribution.h +++ b/infini_train/include/core/distribution.h @@ -42,16 +42,13 @@ namespace cpu { // // Corresponds to PyTorch's cpu::uniform_kernel / cpu::normal_kernel pattern, // simplified to operate on a flat float buffer instead of a TensorIterator. -template -void FillKernel(float *data, int64_t numel, CPUGeneratorImpl *impl, Dist dist) { +template void FillKernel(float *data, int64_t numel, 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); - } + for (int64_t i = 0; i < numel; ++i) { data[i] = dist(engine); } } -} // namespace cpu +} // namespace cpu // ============================== CUDA backend ============================= @@ -68,8 +65,7 @@ namespace cuda { // Locking strategy matches PyTorch: only the offset reservation (PhiloxEngine- // Inputs) touches impl state; the subsequent host-side generation uses purely // local variables and runs outside the impl lock to reduce contention. -template -void FillKernel(float *data, int64_t numel, CUDAGeneratorImpl *impl, Dist dist) { +template void FillKernel(float *data, int64_t numel, CUDAGeneratorImpl *impl, Dist dist) { uint64_t seed, offset; { std::lock_guard lock(impl->mutex_); @@ -80,13 +76,11 @@ void FillKernel(float *data, int64_t numel, CUDAGeneratorImpl *impl, Dist dist) // values coincide (e.g. both backends seeded identically via ManualSeed). constexpr uint64_t kCUDABackendTag = 0xD2511F53D9E08B47ULL; std::mt19937_64 engine(seed ^ ((offset + kCUDABackendTag) * 0x9E3779B97F4A7C15ULL)); - for (int64_t i = 0; i < numel; ++i) { - data[i] = dist(engine); - } + for (int64_t i = 0; i < numel; ++i) { data[i] = dist(engine); } } -} // namespace cuda -#endif // USE_CUDA +} // namespace cuda +#endif // USE_CUDA // ======================== Distribution dispatch ========================== // @@ -102,8 +96,8 @@ void FillKernel(float *data, int64_t numel, CUDAGeneratorImpl *impl, Dist dist) // 2. Dispatch to cpu::FillKernel or cuda::FillKernel // 3. Backend kernel handles locking and element-wise generation template -void FillBuffer(std::vector &buffer, const Device &device, - const std::optional &generator, Dist dist) { +void FillBuffer(std::vector &buffer, const Device &device, const std::optional &generator, + Dist dist) { auto *data = buffer.data(); const auto numel = static_cast(buffer.size()); @@ -124,15 +118,15 @@ void FillBuffer(std::vector &buffer, const Device &device, // Convenience wrappers matching PyTorch's named distribution templates. // uniform_impl_: fills buffer with values from [from, to). -inline void UniformFill(std::vector &buffer, const Device &device, - const std::optional &generator, float from, float to) { +inline void UniformFill(std::vector &buffer, const Device &device, const std::optional &generator, + float from, float to) { FillBuffer(buffer, device, generator, std::uniform_real_distribution(from, to)); } // normal_impl_: fills buffer with values from N(mean, std). -inline void NormalFill(std::vector &buffer, const Device &device, - const std::optional &generator, float mean, float std) { +inline void NormalFill(std::vector &buffer, const Device &device, const std::optional &generator, + float mean, float std) { FillBuffer(buffer, device, generator, std::normal_distribution(mean, std)); } -} // namespace infini_train::core::distribution +} // namespace infini_train::core::distribution diff --git a/infini_train/include/core/generator.h b/infini_train/include/core/generator.h index d85b4375..5c5bb46f 100644 --- a/infini_train/include/core/generator.h +++ b/infini_train/include/core/generator.h @@ -49,7 +49,7 @@ class GeneratorImpl { /* * 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; @@ -91,7 +91,7 @@ class Generator { std::lock_guard lock(impl_->mutex_); return impl_->GetState(); } - void SetState(const std::vector &state) { + void SetState(const std::vector &state) { std::lock_guard lock(impl_->mutex_); impl_->SetState(state); } @@ -126,28 +126,23 @@ 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)); - } + 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); - } + for (int i = 0; i < 8; ++i) { value |= static_cast(p[i]) << (i * 8); } return value; } } // namespace detail -// Global random seed entry point. +// 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) { +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()) @@ -157,16 +152,14 @@ T *GetGeneratorOrDefault(const std::optional &gen, const Device &devi } // Convenience factory — matches PyTorch's at::make_generator(args...). -template -Generator MakeGenerator(Args &&...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) { +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()) diff --git a/infini_train/src/core/cpu_generator.cc b/infini_train/src/core/cpu_generator.cc index c5e70a2c..0f6dac89 100644 --- a/infini_train/src/core/cpu_generator.cc +++ b/infini_train/src/core/cpu_generator.cc @@ -29,7 +29,7 @@ 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_; } @@ -68,8 +68,8 @@ void CPUGeneratorImpl::SetState(const std::vector &state) { 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 << ")"; + 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"; diff --git a/infini_train/src/core/cuda_generator.cc b/infini_train/src/core/cuda_generator.cc index a5634363..cfe3bfce 100644 --- a/infini_train/src/core/cuda_generator.cc +++ b/infini_train/src/core/cuda_generator.cc @@ -2,8 +2,8 @@ #include "infini_train/include/core/cuda_generator.h" -#include #include +#include #include #include #include @@ -67,8 +67,8 @@ void CUDAGeneratorImpl::SetState(const std::vector &state) { 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 << ")"; + 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 diff --git a/infini_train/src/core/distribution.cc b/infini_train/src/core/distribution.cc index 5c0681b8..35e74648 100644 --- a/infini_train/src/core/distribution.cc +++ b/infini_train/src/core/distribution.cc @@ -17,4 +17,4 @@ namespace infini_train::core::distribution { // CHECK_OUT_OF_BOUNDS / check_from_to_in_range) // - Distribution registration for multi-dtype dispatch -} // namespace infini_train::core::distribution +} // namespace infini_train::core::distribution diff --git a/infini_train/src/core/generator.cc b/infini_train/src/core/generator.cc index 58fbdd35..a02b2249 100644 --- a/infini_train/src/core/generator.cc +++ b/infini_train/src/core/generator.cc @@ -48,8 +48,8 @@ Generator &MutableDefaultCUDAGenerator(int8_t 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)); + ? 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; @@ -96,9 +96,7 @@ void ManualSeed(uint64_t 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); - } + for (auto &[index, generator] : detail::MutableDefaultCUDAGenerators()) { generator.SetCurrentSeed(seed); } #endif } diff --git a/infini_train/src/nn/init.cc b/infini_train/src/nn/init.cc index 025f4308..a5dd6277 100644 --- a/infini_train/src/nn/init.cc +++ b/infini_train/src/nn/init.cc @@ -147,9 +147,7 @@ std::shared_ptr Dropout(const std::shared_ptr &tensor, float p, device.IsCPU() ? core::MemcpyKind::kD2D : core::MemcpyKind::kD2H, impl->GetStream(device)); impl->SynchronizeStream(impl->GetStream(device)); - for (int64_t i = 0; i < num_elements; ++i) { - data[i] = (values[i] < p) ? 0.0f : data[i] * scale; - } + for (int64_t i = 0; i < num_elements; ++i) { data[i] = (values[i] < p) ? 0.0f : data[i] * scale; } impl->MemcpyAsync(tensor->DataPtr(), data.data(), num_elements * sizeof(float), device.IsCPU() ? core::MemcpyKind::kD2D : core::MemcpyKind::kH2D, impl->GetStream(device)); diff --git a/tests/generator/test_generator_api.cc b/tests/generator/test_generator_api.cc index b70d0fd4..c2bb1bbd 100644 --- a/tests/generator/test_generator_api.cc +++ b/tests/generator/test_generator_api.cc @@ -24,37 +24,21 @@ class StubGeneratorImpl : public core::GeneratorImpl { explicit StubGeneratorImpl(uint64_t seed = 42) : seed_(seed), offset_(0) {} - void SetCurrentSeed(uint64_t seed) override { - seed_ = seed; - } + void SetCurrentSeed(uint64_t seed) override { seed_ = seed; } - uint64_t CurrentSeed() const override { - return seed_; - } + uint64_t CurrentSeed() const override { return seed_; } - uint64_t Seed() override { - return ++seed_; - } + uint64_t Seed() override { return ++seed_; } - void SetOffset(uint64_t offset) override { - offset_ = offset; - } + void SetOffset(uint64_t offset) override { offset_ = offset; } - uint64_t GetOffset() const override { - return offset_; - } + uint64_t GetOffset() const override { return offset_; } - std::vector GetState() const override { - return std::vector{0xBE, 0xEF}; - } + std::vector GetState() const override { return std::vector{0xBE, 0xEF}; } - void SetState(const std::vector &state) override { - (void)state; - } + void SetState(const std::vector &state) override { (void)state; } - Device GetDevice() const override { - return Device(Device::DeviceType::kCPU, 0); - } + Device GetDevice() const override { return Device(Device::DeviceType::kCPU, 0); } std::shared_ptr Clone() const override { auto cloned = std::make_shared(seed_); @@ -155,7 +139,7 @@ TEST(GeneratorAPITest, BasicAPIForwarding) { 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"); } @@ -168,10 +152,10 @@ TEST(GeneratorAPITest, MakeGeneratorTest) { TEST(GeneratorAPITest, UnsafeGetImplBehavior) { auto gen = core::MakeGenerator(111); - auto* impl = gen.UnsafeGetImpl(); + 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); diff --git a/tests/generator/test_generator_cpu.cc b/tests/generator/test_generator_cpu.cc index 0baba4c1..779922d2 100644 --- a/tests/generator/test_generator_cpu.cc +++ b/tests/generator/test_generator_cpu.cc @@ -1,9 +1,9 @@ // CPU Generator unit tests. -#include "gtest/gtest.h" #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 @@ -22,9 +22,7 @@ std::vector DrawCPU(core::Generator &generator, int64_t count) { std::lock_guard lock(generator.Mutex()); std::vector values(count); auto &engine = impl->Engine(); - for (auto &value : values) { - value = engine(); - } + for (auto &value : values) { value = engine(); } return values; } @@ -125,9 +123,7 @@ TEST(CPUGeneratorTest, RejectsUnsupportedCPUStateVersion) { // Corrupt version field (bytes 8-15) to 99. state[8] = 99; - for (int i = 9; i < 16; ++i) { - state[i] = 0; - } + for (int i = 9; i < 16; ++i) { state[i] = 0; } EXPECT_DEATH(gen.SetState(state), "unsupported version"); } @@ -143,7 +139,8 @@ TEST(CPUGeneratorTest, CPUOnlySupportsZeroOffset) { 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"); + EXPECT_DEATH(core::GetGeneratorOrDefault(gen, Device(Device::DeviceType::kCPU, 0)), + "Generator device type mismatch"); } TEST(CPUGeneratorTest, RejectsMalformedCPUState) { @@ -176,8 +173,6 @@ 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; - } + 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 index 1ee34322..e68c82e0 100644 --- a/tests/generator/test_generator_cuda.cc +++ b/tests/generator/test_generator_cuda.cc @@ -1,8 +1,8 @@ // CUDA Generator unit tests. -#include "gtest/gtest.h" #include "infini_train/include/core/generator.h" #include "infini_train/include/device.h" +#include "gtest/gtest.h" #include #include #include @@ -47,9 +47,7 @@ TEST(CUDAGeneratorTest, RejectsUnsupportedCUDAStateVersion) { // Corrupt version field (bytes 8-15) to 99. state[8] = 99; - for (int i = 9; i < 16; ++i) { - state[i] = 0; - } + for (int i = 9; i < 16; ++i) { state[i] = 0; } EXPECT_DEATH(gen.SetState(state), "unsupported version"); } diff --git a/tests/generator/test_generator_default.cc b/tests/generator/test_generator_default.cc index 9948a689..29db59ca 100644 --- a/tests/generator/test_generator_default.cc +++ b/tests/generator/test_generator_default.cc @@ -1,8 +1,8 @@ // Default generator management and ManualSeed tests. -#include "gtest/gtest.h" #include "infini_train/include/core/generator.h" #include "infini_train/include/device.h" +#include "gtest/gtest.h" #include #include #include @@ -17,9 +17,7 @@ using namespace infini_train; namespace { -Device CPUDevice() { - return Device(Device::DeviceType::kCPU, 0); -} +Device CPUDevice() { return Device(Device::DeviceType::kCPU, 0); } core::Generator MakeCPUGenerator(uint64_t seed) { return core::Generator(std::make_shared(seed)); @@ -30,9 +28,7 @@ std::vector DrawCPU(core::Generator &generator, int64_t count) { std::lock_guard lock(generator.Mutex()); std::vector values(count); auto &engine = impl->Engine(); - for (auto &value : values) { - value = engine(); - } + for (auto &value : values) { value = engine(); } return values; } diff --git a/tests/generator/test_generator_operators.cc b/tests/generator/test_generator_operators.cc index 83285192..68396a57 100644 --- a/tests/generator/test_generator_operators.cc +++ b/tests/generator/test_generator_operators.cc @@ -4,7 +4,6 @@ // KaimingUniform, and Dropout) correctly interact with the Generator mechanism // on both CPU and CUDA backends, leveraging GTest parameterization. -#include "tests/common/test_utils.h" #include "infini_train/include/core/cpu_generator.h" #include "infini_train/include/core/generator.h" #include "infini_train/include/core/runtime/device_guard.h" @@ -12,6 +11,7 @@ #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 @@ -23,9 +23,7 @@ using namespace infini_train; namespace { -Device CPUDevice() { - return Device(Device::DeviceType::kCPU, 0); -} +Device CPUDevice() { return Device(Device::DeviceType::kCPU, 0); } } // namespace class GeneratorOperatorsTest : public infini_train::test::InfiniTrainTest { @@ -300,9 +298,7 @@ TEST_P(GeneratorOperatorsTest, DropoutBoundaryProbabilityZeroKeepsAllElements) { auto gen = MakeDeviceGenerator(3); nn::init::Dropout(tensor, 0.0f, gen); - for (float v : ReadFloats(tensor)) { - EXPECT_FLOAT_EQ(v, 1.0f); - } + for (float v : ReadFloats(tensor)) { EXPECT_FLOAT_EQ(v, 1.0f); } } TEST_P(GeneratorOperatorsTest, DropoutBoundaryProbabilityNearOneDropsMost) { diff --git a/tests/generator/test_generator_thread_safety.cc b/tests/generator/test_generator_thread_safety.cc index 380f9a67..da3bd813 100644 --- a/tests/generator/test_generator_thread_safety.cc +++ b/tests/generator/test_generator_thread_safety.cc @@ -1,7 +1,7 @@ -#include "gtest/gtest.h" -#include "infini_train/include/core/generator.h" #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 @@ -19,9 +19,7 @@ std::vector DrawCPU(core::Generator &gen, int64_t n) { std::lock_guard lock(gen.Mutex()); std::vector values(static_cast(n)); auto &engine = impl->Engine(); - for (auto &v : values) { - v = engine(); - } + for (auto &v : values) { v = engine(); } return values; } @@ -34,19 +32,15 @@ TEST(GeneratorThreadSafety, MultipleThreadsDraw) { auto worker = [&gen]() { for (int i = 0; i < kDrawsPerThread; ++i) { - auto* impl = gen.Get(); + 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(); - } + for (int i = 0; i < kNumThreads; ++i) { threads.emplace_back(worker); } + for (auto &t : threads) { t.join(); } EXPECT_TRUE(gen.Defined()); @@ -66,19 +60,15 @@ TEST(GeneratorThreadSafety, StateConcurrent) { auto state = gen.GetState(); gen.SetState(state); - auto* impl = gen.Get(); + 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(); - } + 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(); @@ -101,12 +91,8 @@ TEST(GeneratorThreadSafety, ManualSeedConcurrent) { }; std::vector threads; - for (int i = 0; i < kNumThreads; ++i) { - threads.emplace_back(worker, static_cast(i) * 1000); - } - for (auto& t : threads) { - t.join(); - } + 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); @@ -126,19 +112,15 @@ TEST(GeneratorThreadSafety, DefaultGeneratorConcurrent) { auto gen = core::detail::DefaultCPUGenerator(); auto state = gen.GetState(); - auto* impl = gen.Get(); + 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(); - } + 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); @@ -162,7 +144,7 @@ TEST(GeneratorThreadSafety, ConcurrentSeedAndDrawConsistency) { } } else { for (int i = 0; i < 100; ++i) { - auto* impl = gen.Get(); + auto *impl = gen.Get(); std::lock_guard lock(gen.Mutex()); (void)impl->Engine()(); } @@ -170,12 +152,8 @@ TEST(GeneratorThreadSafety, ConcurrentSeedAndDrawConsistency) { }; std::vector threads; - for (int i = 0; i < kNumThreads; ++i) { - threads.emplace_back(worker, i); - } - for (auto& t : threads) { - t.join(); - } + 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);