Implements the ThreeFry counter-based PRNG algorithm. Salmon et al. SC 2011. Parallel random numbers: as easy as 1, 2, 3. http://www.thesalmons.org/john/random123/papers/random123sc11.pdf
| 48 | // Salmon et al. SC 2011. Parallel random numbers: as easy as 1, 2, 3. |
| 49 | // http://www.thesalmons.org/john/random123/papers/random123sc11.pdf |
| 50 | ThreeFry2x32State ThreeFry2x32(ThreeFry2x32State input, ThreeFry2x32State key) { |
| 51 | XlaBuilder* builder = input[0].builder(); |
| 52 | key[0] = BitcastConvertType(key[0], U32); |
| 53 | key[1] = BitcastConvertType(key[1], U32); |
| 54 | |
| 55 | // Rotation distances specified by the Threefry2x32 algorithm. |
| 56 | constexpr std::array<int, 8> rotations = {13, 15, 26, 6, 17, 29, 16, 24}; |
| 57 | ThreeFry2x32State x; |
| 58 | |
| 59 | std::array<XlaOp, 3> ks; |
| 60 | // 0x1BD11BDA is a parity constant specified by the ThreeFry2x32 algorithm. |
| 61 | ks[2] = ConstantR0<uint32>(builder, 0x1BD11BDA); |
| 62 | for (int i = 0; i < 2; ++i) { |
| 63 | ks[i] = key[i]; |
| 64 | x[i] = input[i]; |
| 65 | ks[2] = ks[2] ^ key[i]; |
| 66 | } |
| 67 | |
| 68 | x[0] = x[0] + ks[0]; |
| 69 | x[1] = x[1] + ks[1]; |
| 70 | |
| 71 | // Performs a single round of the Threefry2x32 algorithm, with a rotation |
| 72 | // amount 'rotation'. |
| 73 | auto round = [](ThreeFry2x32State v, int rotation) { |
| 74 | v[0] = v[0] + v[1]; |
| 75 | v[1] = RotateLeftU32(v[1], rotation); |
| 76 | v[1] = v[0] ^ v[1]; |
| 77 | return v; |
| 78 | }; |
| 79 | |
| 80 | // There are no known statistical flaws with 13 rounds of Threefry2x32. |
| 81 | // We are conservative and use 20 rounds. |
| 82 | x = round(x, rotations[0]); |
| 83 | x = round(x, rotations[1]); |
| 84 | x = round(x, rotations[2]); |
| 85 | x = round(x, rotations[3]); |
| 86 | x[0] = x[0] + ks[1]; |
| 87 | x[1] = x[1] + ks[2] + ConstantR0<uint32>(builder, 1); |
| 88 | |
| 89 | x = round(x, rotations[4]); |
| 90 | x = round(x, rotations[5]); |
| 91 | x = round(x, rotations[6]); |
| 92 | x = round(x, rotations[7]); |
| 93 | x[0] = x[0] + ks[2]; |
| 94 | x[1] = x[1] + ks[0] + ConstantR0<uint32>(builder, 2); |
| 95 | |
| 96 | x = round(x, rotations[0]); |
| 97 | x = round(x, rotations[1]); |
| 98 | x = round(x, rotations[2]); |
| 99 | x = round(x, rotations[3]); |
| 100 | x[0] = x[0] + ks[0]; |
| 101 | x[1] = x[1] + ks[1] + ConstantR0<uint32>(builder, 3); |
| 102 | |
| 103 | x = round(x, rotations[4]); |
| 104 | x = round(x, rotations[5]); |
| 105 | x = round(x, rotations[6]); |
| 106 | x = round(x, rotations[7]); |
| 107 | x[0] = x[0] + ks[1]; |
no test coverage detected