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/generator.h

@@ -0,0 +1,171 @@
+#pragma once
+
+#include <cstdint>
+#include <memory>
+#include <mutex>
+#include <optional>
+
+#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<uint8_t> GetState() const = 0;
+    virtual void SetState(const std::vector<uint8_t> &state) = 0;
+
+    virtual Device GetDevice() const = 0;
+    virtual std::shared_ptr<GeneratorImpl> 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<GeneratorImpl> impl) : impl_(std::move(impl)) {
+        CHECK(impl_ != nullptr) << "GeneratorImpl with nullptr is not supported";
+    }
+
+    bool Defined() const { return static_cast<bool>(impl_); }
+
+    // Seed control --------------------------------------------------------
+    void SetCurrentSeed(uint64_t seed) {
+        std::lock_guard<std::mutex> lock(impl_->mutex_);
+        impl_->SetCurrentSeed(seed);
+    }
+
+    void ManualSeed(uint64_t seed) { SetCurrentSeed(seed); }
+    uint64_t CurrentSeed() const {
+        std::lock_guard<std::mutex> lock(impl_->mutex_);
+        return impl_->CurrentSeed();
+    }
+
+    uint64_t Seed() {
+        std::lock_guard<std::mutex> lock(impl_->mutex_);
+        return impl_->Seed();
+    }
+
+    void SetOffset(uint64_t offset) {
+        std::lock_guard<std::mutex> lock(impl_->mutex_);
+        impl_->SetOffset(offset);
+    }
+    uint64_t GetOffset() const {
+        std::lock_guard<std::mutex> lock(impl_->mutex_);
+        return impl_->GetOffset();
+    }
+
+    // State control -------------------------------------------------------
+    std::vector<uint8_t> GetState() const {
+        std::lock_guard<std::mutex> lock(impl_->mutex_);
+        return impl_->GetState();
+    }
+    void SetState(const std::vector<uint8_t> &state) {
+        std::lock_guard<std::mutex> 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 <typename T> T *Get() const { return static_cast<T *>(impl_.get()); }
+
+    Generator Clone() const {
+        std::lock_guard<std::mutex> 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<GeneratorImpl> 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<uint8_t> &out, uint64_t value) {
+    for (int i = 0; i < 8; ++i) { out.push_back(static_cast<uint8_t>((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<uint64_t>(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 <typename T> T *GetGeneratorOrDefault(const std::optional<Generator> &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<int>(T::kDeviceType) << " but found "
+        << static_cast<int>(chosen.GetDevice().type());
+    return chosen.Get<T>();
+}
+
+// Convenience factory — matches PyTorch's at::make_generator<Impl>(args...).
+template <class Impl, class... Args> Generator MakeGenerator(Args &&...args) {
+    return Generator(std::make_shared<Impl>(std::forward<Args>(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 <typename T> T *CheckGenerator(std::optional<Generator> 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<int>(T::kDeviceType) << " but found "
+        << static_cast<int>(gen->GetDevice().type());
+    return gen->Get<T>();
+}
+
+} // namespace infini_train::core

infini_train/include/nn/init.h

@@ -2,9 +2,9 @@
 
 #include <memory>
 #include <optional>
-#include <random>
 #include <utility>
 
+#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<Tensor> Normal(const std::shared_ptr<Tensor> &tensor, float mean = 0.0, float std = 1.0,
-                               std::optional<std::mt19937> generator = std::nullopt);
+                               std::optional<core::Generator> generator = std::nullopt);
 
 std::pair<int64_t, int64_t> CalculateFanInAndFanOut(const std::shared_ptr<Tensor> &tensor);
 
@@ -42,10 +42,13 @@ enum class NonLinearityType : int8_t {
 std::shared_ptr<Tensor> KaimingUniform(const std::shared_ptr<Tensor> &tensor, float a = 0.0f,
                                        KaimingMode mode = KaimingMode::kFanIn,
                                        NonLinearityType non_linearity = NonLinearityType::kLeakyReLU,
-                                       std::optional<std::mt19937> generator = std::nullopt);
+                                       std::optional<core::Generator> generator = std::nullopt);
 
 std::shared_ptr<Tensor> Uniform(const std::shared_ptr<Tensor> &tensor, float a = 0.0f, float b = 1.0f,
-                                std::optional<std::mt19937> generator = std::nullopt);
+                                std::optional<core::Generator> generator = std::nullopt);
+
+std::shared_ptr<Tensor> Dropout(const std::shared_ptr<Tensor> &tensor, float p = 0.5f,
+                                std::optional<core::Generator> generator = std::nullopt);
 
 std::shared_ptr<Tensor> Ones(const std::shared_ptr<Tensor> &tensor);
 

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<Tensor> {
 
     // distribution
     std::shared_ptr<Tensor> Uniform(float from = 0.0f, float to = 1.0f,
-                                    std::optional<std::mt19937> generator = std::nullopt);
+                                    std::optional<core::Generator> generator = std::nullopt);
 
     std::shared_ptr<Tensor> Matmul(const std::shared_ptr<Tensor> &other);
     std::shared_ptr<Tensor> Outer(const std::shared_ptr<Tensor> &other);