Determines whether to use the composite or specialized CPU kernel. When the total size of the tensor is larger than the cache size and the batch size is large compared to the smallest matrix dimension, then the composite implementation is inefficient since it has to read the entire
(fast, tensor_shape)
| 229 | |
| 230 | # pylint: disable=long-lambda |
| 231 | def _use_composite_impl(fast, tensor_shape): |
| 232 | """Determines whether to use the composite or specialized CPU kernel. |
| 233 | |
| 234 | When the total size of the tensor is larger than the cache size and the |
| 235 | batch size is large compared to the smallest matrix dimension, then the |
| 236 | composite implementation is inefficient since it has to read the entire |
| 237 | tensor from memory multiple times. In this case we fall back to the |
| 238 | original CPU kernel, which does all the computational steps on each |
| 239 | matrix separately. |
| 240 | |
| 241 | Only fast mode is supported by the composite impl, so `False` is returned |
| 242 | if `fast` is `False`. |
| 243 | |
| 244 | Args: |
| 245 | fast: bool indicating if fast mode in the solver was requested. |
| 246 | tensor_shape: The shape of the tensor. |
| 247 | |
| 248 | Returns: |
| 249 | True if the composite impl should be used. False otherwise. |
| 250 | """ |
| 251 | if fast is False: |
| 252 | return False |
| 253 | batch_shape = tensor_shape[:-2] |
| 254 | matrix_shape = tensor_shape[-2:] |
| 255 | if not tensor_shape.is_fully_defined(): |
| 256 | return True |
| 257 | tensor_size = tensor_shape.num_elements() * matrix.dtype.size |
| 258 | is_io_bound = batch_shape.num_elements() > np.min(matrix_shape) |
| 259 | L2_CACHE_SIZE_GUESSTIMATE = 256000 |
| 260 | if tensor_size > L2_CACHE_SIZE_GUESSTIMATE and is_io_bound: |
| 261 | return False |
| 262 | else: |
| 263 | return True |
| 264 | |
| 265 | def _overdetermined(matrix, rhs, l2_regularizer): |
| 266 | """Computes (A^H*A + l2_regularizer)^{-1} * A^H * rhs.""" |
no test coverage detected