2026-spring-littleotherut-T2-2-1

infini_train/include/core/cpu_generator.h

@@ -0,0 +1,57 @@
+#pragma once
+
+#include <cstdint>
+#include <memory>
+#include <random>
+#include <vector>
+
+#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<uint8_t> GetState() const override;
+    void SetState(const std::vector<uint8_t> &state) override;
+
+    Device GetDevice() const override;
+    std::shared_ptr<GeneratorImpl> 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

infini_train/include/core/cuda_generator.h

@@ -0,0 +1,63 @@
+#pragma once
+
+#include <cstdint>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#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<uint8_t> GetState() const override;
+    void SetState(const std::vector<uint8_t> &state) override;
+
+    Device GetDevice() const override;
+    std::shared_ptr<GeneratorImpl> 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<uint64_t, uint64_t> 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

infini_train/include/core/distribution.h

@@ -0,0 +1,132 @@
+#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 <cstdint>
+#include <mutex>
+#include <optional>
+#include <random>
+#include <vector>
+
+#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 <typename Dist> void FillKernel(float *data, int64_t numel, CPUGeneratorImpl *impl, Dist dist) {
+    std::lock_guard<std::mutex> 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 <typename Dist> void FillKernel(float *data, int64_t numel, CUDAGeneratorImpl *impl, Dist dist) {
+    uint64_t seed, offset;
+    {
+        std::lock_guard<std::mutex> lock(impl->mutex_);
+        std::tie(seed, offset) = impl->PhiloxEngineInputs(static_cast<uint64_t>(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 <typename Dist>
+void FillBuffer(std::vector<float> &buffer, const Device &device, const std::optional<Generator> &generator,
+                Dist dist) {
+    auto *data = buffer.data();
+    const auto numel = static_cast<int64_t>(buffer.size());
+
+    if (device.IsCPU()) {
+        auto *impl = GetGeneratorOrDefault<CPUGeneratorImpl>(generator, device);
+        cpu::FillKernel(data, numel, impl, dist);
+        return;
+    }
+
+#ifdef USE_CUDA
+    auto *impl = GetGeneratorOrDefault<CUDAGeneratorImpl>(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<float> &buffer, const Device &device, const std::optional<Generator> &generator,
+                        float from, float to) {
+    FillBuffer(buffer, device, generator, std::uniform_real_distribution<float>(from, to));
+}
+
+// normal_impl_: fills buffer with values from N(mean, std).
+inline void NormalFill(std::vector<float> &buffer, const Device &device, const std::optional<Generator> &generator,
+                       float mean, float std) {
+    FillBuffer(buffer, device, generator, std::normal_distribution<float>(mean, std));
+}
+
+} // namespace infini_train::core::distribution