crc.py (11651B)
1 # 2 # SPDX-FileCopyrightText: Copyright (c) 2021-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 # SPDX-License-Identifier: Apache-2.0 4 # 5 """Layers for cyclic redundancy checks (CRC) and utility functions""" 6 7 import numpy as np 8 import tensorflow as tf 9 from tensorflow.keras.layers import Layer 10 from sionna.fec.utils import int_mod_2 11 12 class CRCEncoder(Layer): 13 """CRCEncoder(crc_degree, output_dtype=tf.float32, **kwargs) 14 15 Adds cyclic redundancy check to input sequence. 16 17 The CRC polynomials from Sec. 5.1 in [3GPPTS38212_CRC]_ are available: 18 `{CRC24A, CRC24B, CRC24C, CRC16, CRC11, CRC6}`. 19 20 The class inherits from the Keras layer class and can be used as layer in a 21 Keras model. 22 23 Parameters 24 ---------- 25 crc_degree: str 26 Defining the CRC polynomial to be used. Can be any value from 27 `{CRC24A, CRC24B, CRC24C, CRC16, CRC11, CRC6}`. 28 29 dtype: tf.DType 30 Defaults to `tf.float32`. Defines the output dtype. 31 32 Input 33 ----- 34 inputs : [...,k], tf.float32 35 2+D tensor of arbitrary shape where the last dimension is 36 `[...,k]`. Must have at least rank two. 37 38 Output 39 ------ 40 x_crc : [...,k+crc_degree], tf.float32 41 2+D tensor containing CRC encoded bits of same shape as 42 ``inputs`` except the last dimension changes to 43 `[...,k+crc_degree]`. 44 45 Raises 46 ------ 47 AssertionError 48 If ``crc_degree`` is not `str`. 49 50 ValueError 51 If requested CRC polynomial is not available in [3GPPTS38212_CRC]_. 52 53 InvalidArgumentError 54 When rank(``inputs``)<2. 55 56 Note 57 ---- 58 For performance enhancements, we implement a generator-matrix based 59 implementation for fixed `k` instead of the more common shift 60 register-based operations. Thus, the encoder need to trigger an 61 (internal) rebuild if `k` changes. 62 63 """ 64 65 def __init__(self, crc_degree, dtype=tf.float32, **kwargs): 66 67 super().__init__(dtype=dtype, **kwargs) 68 69 assert isinstance(crc_degree, str), "crc_degree must be str" 70 self._crc_degree = crc_degree 71 72 # init 5G CRC polynomial 73 self._crc_pol, self._crc_length = self._select_crc_pol(self._crc_degree) 74 75 self._k = None 76 self._n = None 77 78 ######################################### 79 # Public methods and properties 80 ######################################### 81 82 @property 83 def crc_degree(self): 84 """CRC degree as string.""" 85 return self._crc_degree 86 87 @property 88 def crc_length(self): 89 """Length of CRC. Equals number of CRC parity bits.""" 90 return self._crc_length 91 92 @property 93 def crc_pol(self): 94 """CRC polynomial in binary representation.""" 95 return self._crc_pol 96 97 @property 98 def k(self): 99 """Number of information bits per codeword.""" 100 return self._k 101 102 @property 103 def n(self): 104 """Number of codeword bits.""" 105 return self._n 106 107 ######################### 108 # Utility methods 109 ######################### 110 111 def _select_crc_pol(self, crc_degree): 112 """Select 5G CRC polynomial according to Sec. 5.1 [3GPPTS38212_CRC]_.""" 113 if crc_degree=="CRC24A": 114 crc_length = 24 115 crc_coeffs = [24, 23, 18, 17, 14, 11, 10, 7, 6, 5, 4, 3, 1, 0] 116 elif crc_degree=="CRC24B": 117 crc_length = 24 118 crc_coeffs = [24, 23, 6, 5, 1, 0] 119 elif crc_degree=="CRC24C": 120 crc_length = 24 121 crc_coeffs = [24, 23, 21, 20, 17, 15, 13, 12, 8, 4, 2, 1, 0] 122 elif crc_degree=="CRC16": 123 crc_length = 16 124 crc_coeffs = [16, 12, 5, 0] 125 elif crc_degree=="CRC11": 126 crc_length = 11 127 crc_coeffs = [11, 10, 9, 5, 0] 128 elif crc_degree=="CRC6": 129 crc_length = 6 130 crc_coeffs = [6, 5, 0] 131 else: 132 raise ValueError("Invalid CRC Polynomial") 133 134 crc_pol = np.zeros(crc_length+1) 135 for c in crc_coeffs: 136 crc_pol[c] = 1 137 138 # invert array (MSB instead of LSB) 139 crc_pol_inv = np.zeros(crc_length+1) 140 for i in range(crc_length+1): 141 crc_pol_inv[crc_length-i] = crc_pol[i] 142 143 return crc_pol_inv.astype(int), crc_length 144 145 def _gen_crc_mat(self, k, pol_crc): 146 """ Build (dense) generator matrix for CRC parity bits. 147 148 The principle idea is to treat the CRC as systematic linear code, i.e., 149 the generator matrix can be composed out of ``k`` linear independent 150 (valid) codewords. For this, we CRC encode all ``k`` unit-vectors 151 `[0,...1,...,0]` and build the generator matrix. 152 To avoid `O(k^2)` complexity, we start with the last unit vector 153 given as `[0,...,0,1]` and can generate the result for next vector 154 `[0,...,1,0]` via another polynom division of the remainder from the 155 previous result. This allows to successively build the generator matrix 156 at linear complexity `O(k)`. 157 """ 158 crc_length = len(pol_crc) - 1 159 g_mat = np.zeros([k, crc_length]) 160 161 x_crc = np.zeros(crc_length).astype(int) 162 x_crc[0] = 1 163 for i in range(k): 164 # shift by one position 165 x_crc = np.concatenate([x_crc, [0]]) 166 if x_crc[0]==1: 167 x_crc = np.bitwise_xor(x_crc, pol_crc) 168 x_crc = x_crc[1:] 169 g_mat[k-i-1,:] = x_crc 170 171 return g_mat 172 173 ######################### 174 # Keras layer functions 175 ######################### 176 177 def build(self, input_shape): 178 """Build the generator matrix 179 180 The CRC is always added to the last dimension of the input. 181 """ 182 k = input_shape[-1] # we perform the CRC check on the last dimension 183 assert k is not None, "Shape of last dimension cannot be None." 184 g_mat_crc = self._gen_crc_mat(k, self.crc_pol) 185 self._g_mat_crc = tf.constant(g_mat_crc, dtype=tf.float32) 186 187 self._k = k 188 self._n = k + g_mat_crc.shape[1] 189 190 def call(self, inputs): 191 """cyclic redundancy check function. 192 193 This function add the CRC parity bits to ``inputs`` and returns the 194 result of the CRC validation. 195 196 Args: 197 inputs (tf.float32): Tensor of arbitrary shape `[...,k]`. 198 Must have at least rank two. 199 200 Returns: 201 `tf.float32`: CRC encoded bits ``x_crc``of shape 202 `[...,k+crc_degree]`. 203 204 Raises: 205 InvalidArgumentError: When rank(``x``)<2. 206 207 """ 208 209 # assert rank must be two 210 tf.debugging.assert_greater(tf.rank(inputs), 1) 211 212 # re-init if shape has changed, update generator matrix 213 if inputs.shape[-1] != self._g_mat_crc.shape[0]: 214 self.build(inputs.shape) 215 216 # note: as the code is systematic, we only encode the crc positions 217 # this the generator matrix is non-sparse and a "full" matrix 218 # multiplication is probably the fastest implementation. 219 220 x_exp = tf.expand_dims(inputs, axis=-2) # row vector of shape 1xk 221 222 # tf.matmul onl supports floats (and int32 but not uint8 etc.) 223 x_exp32 = tf.cast(x_exp, tf.float32) 224 x_crc = tf.matmul(x_exp32, self._g_mat_crc) # calculate crc bits 225 226 # take modulo 2 of x_crc (bitwise operations instead of tf.mod) 227 228 # cast to tf.int64 first as TF 2.15 has an XLA bug with casting directly 229 # to tf.int32 230 x_crc = tf.cast(x_crc, dtype=tf.int64) 231 232 x_crc = int_mod_2(x_crc) 233 x_crc = tf.cast(x_crc, dtype=self.dtype) 234 235 x_conc = tf.concat([x_exp, x_crc], -1) 236 x_out = tf.squeeze(x_conc, axis=-2) 237 238 return x_out 239 240 241 class CRCDecoder(Layer): 242 """CRCDecoder(crc_encoder, dtype=None, **kwargs) 243 244 Allows cyclic redundancy check verification and removes parity-bits. 245 246 The CRC polynomials from Sec. 5.1 in [3GPPTS38212_CRC]_ are available: 247 `{CRC24A, CRC24B, CRC24C, CRC16, CRC11, CRC6}`. 248 249 The class inherits from the Keras layer class and can be used as layer in a 250 Keras model. 251 252 Parameters 253 ---------- 254 crc_encoder: CRCEncoder 255 An instance of :class:`~sionna.fec.crc.CRCEncoder` to which the 256 CRCDecoder is associated. 257 258 dtype: tf.DType 259 Defaults to `None`. Defines the datatype for internal calculations 260 and the output dtype. If no explicit dtype is provided the dtype 261 from the associated interleaver is used. 262 263 Input 264 ----- 265 inputs: [...,k+crc_degree], tf.float32 266 2+D Tensor containing the CRC encoded bits (i.e., the last 267 `crc_degree` bits are parity bits). Must have at least rank two. 268 269 Output 270 ------ 271 (x, crc_valid): 272 Tuple: 273 274 x : [...,k], tf.float32 275 2+D tensor containing the information bit sequence without CRC 276 parity bits. 277 278 crc_valid : [...,1], tf.bool 279 2+D tensor containing the result of the CRC per codeword. 280 281 Raises 282 ------ 283 AssertionError 284 If ``crc_encoder`` is not `CRCEncoder`. 285 286 InvalidArgumentError 287 When rank(``x``)<2. 288 """ 289 290 def __init__(self, 291 crc_encoder, 292 dtype=tf.float32, 293 **kwargs): 294 295 assert isinstance(crc_encoder, CRCEncoder), \ 296 "crc_encoder must be an instance of CRCEncoder." 297 self._encoder = crc_encoder 298 299 # if dtype is None, use same dtype as associated encoder 300 if dtype is None: 301 dtype = self._encoder.dtype 302 303 super().__init__(dtype=dtype, **kwargs) 304 305 ######################################### 306 # Public methods and properties 307 ######################################### 308 309 @property 310 def crc_degree(self): 311 """CRC degree as string.""" 312 return self._encoder.crc_degree 313 314 @property 315 def encoder(self): 316 """CRC Encoder used for internal validation.""" 317 return self._encoder 318 319 ######################### 320 # Keras layer functions 321 ######################### 322 323 def build(self, input_shape): 324 """Nothing to build.""" 325 pass 326 327 def call(self, inputs): 328 """cyclic redundancy check verification 329 330 This function verifies the CRC of ``inputs``. Returns the result of 331 the CRC validation and removes parity bits from ``inputs``. 332 333 Args: 334 inputs (tf.float32): Tensor of arbitrary shape `[...,k+crc_degree]`. 335 Must have at least rank two. 336 337 Returns: 338 List(`tf.float32`, `tf.bool`): ``[x, crc_valid]`` list of the 339 information bits ``x`` and the result of the parity check 340 validation ``crc_valid`` of each codeword, where ``x`` has shape 341 `[...,k]` and ``crc_valid`` has shape `[...,1]`. 342 343 Raises: 344 InvalidArgumentError: When rank(``inputs``)<2. 345 346 """ 347 348 # assert rank must be two 349 tf.debugging.assert_greater(tf.rank(inputs), 1) 350 351 # last dim must be at least crc_bits long 352 tf.debugging.assert_greater_equal(tf.shape(inputs)[-1], 353 self._encoder.crc_length) 354 355 # re-encode information bits of x and verify that CRC bits are correct 356 357 x_info = inputs[...,0:-self._encoder.crc_length] 358 x_parity = self._encoder(inputs)[...,-self._encoder.crc_length:] 359 360 # return if x fulfils the CRC 361 crc_check = tf.reduce_sum(x_parity, axis=-1, keepdims=True) 362 crc_check = tf.where(crc_check>0, False, True) 363 364 return [x_info, crc_check]