(weight)
| 15 | |
| 16 | |
| 17 | def permutate_weight_fastest(weight): |
| 18 | wmma_n = 16 |
| 19 | wmma_k = 32 |
| 20 | N = weight.shape[0] |
| 21 | K = weight.shape[1] |
| 22 | |
| 23 | # Create a lookup table for the permutation |
| 24 | mapping = np.zeros((wmma_n, wmma_k, 2), dtype=int) |
| 25 | for ii in range(wmma_n): |
| 26 | for jj in range(wmma_k): |
| 27 | mapping[ii, jj] = B_global_16x32_to_shared_load_16x32_layout(ii, jj) |
| 28 | |
| 29 | # Reshape weight for the final format |
| 30 | permutated_weight = np.zeros((N // wmma_n, K // wmma_k, wmma_n, wmma_k), dtype="int8") |
| 31 | |
| 32 | # Use advanced indexing for the entire operation |
| 33 | i_indices = np.arange(N // wmma_n)[:, np.newaxis, np.newaxis, np.newaxis] |
| 34 | j_indices = np.arange(K // wmma_k)[np.newaxis, :, np.newaxis, np.newaxis] |
| 35 | |
| 36 | # Create the source indices |
| 37 | src_i = i_indices * wmma_n + mapping[:, :, 0] |
| 38 | src_j = j_indices * wmma_k + mapping[:, :, 1] |
| 39 | |
| 40 | # Extract and reshape in one go |
| 41 | permutated_weight = weight[src_i, src_j] |
| 42 | |
| 43 | return permutated_weight |
| 44 | |
| 45 | |
| 46 | def compress_int2_to_int8(int2_weight): |
no test coverage detected