1/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16// This file implements random_index_shuffle() by using a simple block chiper
17// for pseudorandom permutations.
18// This idea is described as cycle-walking cipher in
19// https://www.cs.ucdavis.edu/~rogaway/papers/subset.pdf
20//
21// We use the Simon block cipher described in
22// https://eprint.iacr.org/2013/404
23// and following recommendations in
24// https://nsacyber.github.io/simon-speck/implementations/ImplementationGuide1.1.pdf.
25// However we use a single fixed key size and support arbtitary block sizes.
26// Further we fixed the number of rounds in the Feistel structuro to be always
27// 4. This reduces the computational cost and still gives good shuffle behavior.
28//
29// Warning: Given the modifications descripted above this implementation should
30// not be used for application that require cryptograhic secure RNGs.
31
32#include "tensorflow/core/kernels/random_index_shuffle.h"
33
34#include <assert.h>
35
36#include <algorithm>
37#include <array>
38#include <bitset>
39#include <cmath>
40
41#include "tensorflow/core/platform/types.h"
42
43namespace tensorflow {
44namespace random {
45// Some of the macros below require a minimum word size of 8
46// (2 word size = block size).
47// Smaller block sizes might give poor results in terms of randomness.
48constexpr int kMinBlockSize = 16;
49
50namespace impl {
51
52#define ROTL(x, r, W) (((x) << (r)) | (x >> (W - (r))))
53#define ROTR(x, r, W) (((x) >> (r)) | ((x) << (W - (r))))
54#define SIMON_F(x, W) ((ROTL(x, 1, W) & ROTL(x, 8, W) ^ ROTL(x, 2, W)))
55#define SIMON_Rx2(x, y, k1, k2, W) \
56 (y ^= SIMON_F(x, W), y ^= k1, x ^= SIMON_F(y, W), x ^= k2)
57
58// Returns the keys per round for a Simon cipher.
59// This variant uses std::bitset and can generate keys with any number of bits.
60// This should not be used to encrypt data. It's not secure. We only use it
61// to generate pseudorandom permutations.
62template <int W>
63std::vector<std::bitset<W>> simon_key_schedule(
64 const std::array<uint32_t, 3>& key, const int32_t rounds) {
65 // Required by ROTR/ROTL
66 static_assert(W >= 8, "Minimum word size is 8 bits.");
67 const auto c = std::bitset<W>(0xfffffffc);
68 auto z = std::bitset<W>(0x7369f885192c0ef5LL);
69 std::vector<std::bitset<W>> rk({key[0], key[1], key[2]});
70 rk.reserve(rounds);
71 for (int i = 3; i < rounds; i++) {
72 rk.push_back(c ^ (z & std::bitset<W>(1)) ^ rk[i - 3] ^
73 ROTR(rk[i - 1], 3, W) ^ ROTR(rk[i - 1], 4, W));
74 z >>= 1;
75 }
76 return rk;
77}
78
79// Encrypts the given value using the Simon chipher.
80// This is should not be used to encrypt data. It's not secure. We only use it
81// to generate pseudorandom permutations.
82template <int W>
83uint64_t simon_encrypt(const uint64_t value,
84 const std::vector<std::bitset<W>>& round_keys) {
85 // Required by ROTR/ROTL
86 static_assert(W >= 8, "Minimum word size is 8 bits.");
87 std::bitset<W> left(value >> W);
88 std::bitset<W> right(value);
89 for (int i = 0; i < round_keys.size();) {
90 SIMON_Rx2(right, left, round_keys[i++], round_keys[i++], W);
91 }
92 return (left.to_ullong() << W) | right.to_ullong();
93}
94
95// In the original implementation the number of rounds depends on the block
96// size, key size and key words. For our purposes of random shuffle a 4 to 8
97// rounds is enough.
98// W = word size
99// B = 2 * W = block size
100template <int B>
101uint64_t index_shuffle(const uint64_t index, const std::array<uint32_t, 3>& key,
102 const uint64_t max_index, const int32_t rounds) {
103 const auto round_keys = simon_key_schedule<B / 2>(key, rounds);
104 uint64_t new_index = index;
105 while (true) {
106 new_index = simon_encrypt<B / 2>(new_index, round_keys);
107 if (new_index <= max_index) {
108 return new_index;
109 }
110 }
111}
112
113#undef ROTL
114#undef ROTR
115#undef SIMON_F
116#undef SIMON_RxC
117
118} // namespace impl
119
120uint64_t index_shuffle(const uint64_t index, const std::array<uint32_t, 3>& key,
121 const uint64_t max_index, const int32_t rounds) {
122 // Block size must be large enough to represent max_index and even (since
123 // word size is half of it). We force at least 16 bits as minimum block size
124 // since we observed pattern in the permutations below.
125 int block_size = static_cast<int>(std::ceil(std::log2(max_index)));
126 block_size = std::max(block_size + block_size % 2, kMinBlockSize);
127 assert(block_size > 0 && block_size % 2 == 0 && block_size <= 64);
128 // At least 4 rounds and number of rounds must be even.
129 assert(rounds >= 4 && rounds % 2 == 0);
130#define HANDLE_BLOCK_SIZE(B) \
131 case B: \
132 return impl::index_shuffle<B>(index, key, max_index, rounds);
133
134 switch (block_size) {
135 HANDLE_BLOCK_SIZE(16);
136 HANDLE_BLOCK_SIZE(18);
137 HANDLE_BLOCK_SIZE(20);
138 HANDLE_BLOCK_SIZE(22);
139 HANDLE_BLOCK_SIZE(24);
140 HANDLE_BLOCK_SIZE(26);
141 HANDLE_BLOCK_SIZE(28);
142 HANDLE_BLOCK_SIZE(30);
143 HANDLE_BLOCK_SIZE(32);
144 HANDLE_BLOCK_SIZE(34);
145 HANDLE_BLOCK_SIZE(36);
146 HANDLE_BLOCK_SIZE(38);
147 HANDLE_BLOCK_SIZE(40);
148 HANDLE_BLOCK_SIZE(42);
149 HANDLE_BLOCK_SIZE(44);
150 HANDLE_BLOCK_SIZE(46);
151 HANDLE_BLOCK_SIZE(48);
152 HANDLE_BLOCK_SIZE(50);
153 HANDLE_BLOCK_SIZE(52);
154 HANDLE_BLOCK_SIZE(54);
155 HANDLE_BLOCK_SIZE(56);
156 HANDLE_BLOCK_SIZE(58);
157 HANDLE_BLOCK_SIZE(60);
158 HANDLE_BLOCK_SIZE(62);
159 default:
160 return impl::index_shuffle<64>(index, key, max_index, rounds);
161 }
162#undef HANDLE_BLOCK_SIZE
163}
164
165} // namespace random
166} // namespace tensorflow
167