1 #include <gtest/gtest.h>
2
3 #include <ATen/core/PhiloxRNGEngine.h>
4
5 #include <cstdint>
6 #include <iostream>
7 namespace torch {
8 namespace aot_inductor {
9
randint64_cpu(uint32_t seed,uint32_t offset,int64_t low,int64_t high)10 int64_t randint64_cpu(
11 uint32_t seed,
12 uint32_t offset,
13 int64_t low,
14 int64_t high) {
15 auto gen = at::Philox4_32(seed, 0, offset);
16 uint64_t r0 = gen();
17 uint64_t r1 = gen();
18 uint64_t result = r0 | (r1 << 32);
19 return static_cast<int64_t>(result % (high - low)) + low;
20 }
21
TEST(TestRand,TestRandn)22 TEST(TestRand, TestRandn) {
23 at::Philox4_32 engine_1(1, 0, 0);
24 float a = engine_1.randn(10);
25 at::Philox4_32 engine_2(1, 0, 0);
26 float b = engine_2.randn(10);
27
28 EXPECT_EQ(a, b);
29 }
30
TEST(TestRand,TestRandint64)31 TEST(TestRand, TestRandint64) {
32 int64_t a = randint64_cpu(0xffffffff, 100, 0, INT64_MAX);
33 int64_t b = randint64_cpu(0xffffffff, 100, 0, INT64_MAX);
34
35 EXPECT_EQ(a, b);
36 }
37
38 } // namespace aot_inductor
39 } // namespace torch
40