[CPU] RandomUniform seed alignment

src/plugins/intel_cpu/src/nodes/kernels/x64/random_uniform.hpp

@@ -11,12 +11,18 @@
 namespace ov {
 namespace intel_cpu {
 namespace kernel {
+namespace random_uniform {
 
-struct RandomUniformCompileParams {
+struct PhiloxGeneratorCompileParams {
     element::Type out_data_type = element::f32;
 };
 
-struct RandomUniformCallArgs {
+struct MersenneTwisterGeneratorCompileParams {
+    element::Type out_data_type = element::f32;
+    bool optimized = false;
+};
+
+struct PhiloxGeneratorCallArgs {
     void* dst_ptr;
     const void* key_ptr;
     const void* counter_ptr;
@@ -26,12 +32,23 @@ struct RandomUniformCallArgs {
     uint64_t work_amount = 0lu;
 };
 
+struct MersenneTwisterGeneratorCallArgs {
+    void* dst_ptr;
+    void* state_ptr;
+    const void* min_ptr;
+    const void* range_ptr;
+    uint64_t output_idx = 0;
+    uint64_t max_output_idx = 0;
+    uint64_t state_accesses_count = 0lu;
+    int64_t elements_to_generate = 0lu;
+};
+
 template <dnnl::impl::cpu::x64::cpu_isa_t isa>
-class RandomUniform : public JitKernel<RandomUniformCompileParams, RandomUniformCallArgs> {
+class PhiloxGenerator : public JitKernel<PhiloxGeneratorCompileParams, PhiloxGeneratorCallArgs> {
 public:
-    DECLARE_CPU_JIT_AUX_FUNCTIONS(RandomUniform)
+    DECLARE_CPU_JIT_AUX_FUNCTIONS(PhiloxGenerator)
 
-    explicit RandomUniform(const RandomUniformCompileParams& jcp);
+    explicit PhiloxGenerator(const PhiloxGeneratorCompileParams& jcp);
 
     void generate() override;
 
@@ -102,6 +119,71 @@ class RandomUniform : public JitKernel<RandomUniformCompileParams, RandomUniform
     static constexpr uint64_t STATISTIC_MAXIMIZING_MULTIPLIER_COUNTER = 0xCD9E8D57;
 };
 
+template <dnnl::impl::cpu::x64::cpu_isa_t isa>
+class MersenneTwisterGenerator
+    : public JitKernel<MersenneTwisterGeneratorCompileParams, MersenneTwisterGeneratorCallArgs> {
+public:
+    DECLARE_CPU_JIT_AUX_FUNCTIONS(MersenneTwisterGenerator)
+
+    explicit MersenneTwisterGenerator(const MersenneTwisterGeneratorCompileParams& jcp);
+
+    void generate() override;
+
+private:
+    using Vmm = typename dnnl::impl::utils::conditional3<isa == dnnl::impl::cpu::x64::avx512_core,
+                                                         Xbyak::Zmm,
+                                                         isa == dnnl::impl::cpu::x64::sse41,
+                                                         Xbyak::Xmm,
+                                                         Xbyak::Ymm>::type;
+    using Vmask = typename dnnl::impl::utils::conditional3<isa == dnnl::impl::cpu::x64::avx512_core,
+                                                           Xbyak::Opmask,
+                                                           isa == dnnl::impl::cpu::x64::sse41,
+                                                           Xbyak::Xmm,
+                                                           Xbyak::Ymm>::type;
+
+    RegistersPool::Reg<Xbyak::Reg64> r64_dst;
+    RegistersPool::Reg<Xbyak::Reg64> r64_state;
+    RegistersPool::Reg<Xbyak::Reg64> r64_state_accesses_count;
+    RegistersPool::Reg<Xbyak::Reg64> r64_elements_to_generate;
+    RegistersPool::Reg<Xbyak::Reg64> r64_storage_capacity;
+    RegistersPool::Reg<Xbyak::Reg64> r64_output_idx;
+    RegistersPool::Reg<Xbyak::Reg64> r64_max_output_idx;
+
+    const Xbyak::Reg64 r64_params = Xbyak::Reg64(dnnl::impl::cpu::x64::abi_param_regs[0]);
+
+    // Vector registers for input storage.
+    RegistersPool::Reg<Vmm> v_dst;
+    RegistersPool::Reg<Vmm> v_state;
+    RegistersPool::Reg<Vmm> v_min;
+    RegistersPool::Reg<Vmm> v_range;
+
+    // Vector registers for generation.
+    RegistersPool::Reg<Vmm> v_const_1;  // for MT_CONST_1
+    RegistersPool::Reg<Vmm> v_const_2;  // for MT_CONST_2
+
+    // Vector registers for conversion.
+    RegistersPool::Reg<Vmm> v_mask;
+    RegistersPool::Reg<Vmm> v_divisor;
+
+    // Output vector register
+    RegistersPool::Reg<Vmm> v_result;
+
+    void initVectors();
+
+    void process();
+
+    void generateRandomNumbers();
+
+    void convertToOutputTypeMersenne();
+
+    void storeResults();
+
+    // Mersenne Twister constants
+    static constexpr uint32_t MT_CONST_1 = 0x9D2C5680;
+    static constexpr uint32_t MT_CONST_2 = 0xEFC60000;
+};
+
+}  // namespace random_uniform
 }  // namespace kernel
 }  // namespace intel_cpu
 }  // namespace ov