compress the tensor on cpu # Convert to 4-bit metadata value # The metadata value represents which 2 elements are non-zero # 0x4: [1,1,0,0] - first two elements are non-zero # 0x8: [1,0,1,0] - first and third elements are non-zero # 0xC: [1,0,0,1] - f
(self, a, a_compressed, meta)
| 88 | return self.__compress_on_cuda(a, a_compressed, meta) |
| 89 | |
| 90 | def __compress_on_cpu(self, a, a_compressed, meta): |
| 91 | """ |
| 92 | compress the tensor on cpu |
| 93 | # Convert to 4-bit metadata value |
| 94 | # The metadata value represents which 2 elements are non-zero |
| 95 | # 0x4: [1,1,0,0] - first two elements are non-zero |
| 96 | # 0x8: [1,0,1,0] - first and third elements are non-zero |
| 97 | # 0xC: [1,0,0,1] - first and fourth elements are non-zero |
| 98 | # 0x9: [0,1,1,0] - second and third elements are non-zero |
| 99 | # 0xD: [0,1,0,1] - second and fourth elements are non-zero |
| 100 | # 0xE: [0,0,1,1] - third and fourth elements are non-zero |
| 101 | # special case: |
| 102 | # [0,0,0,0] == [0,0,1,1] |
| 103 | # [1,0,0,0] == [1,0,0,1] |
| 104 | # [0,1,0,0] == [0,1,0,1] |
| 105 | # [0,0,1,0] == [0,0,1,1] |
| 106 | # [0,0,0,1] == [0,0,1,1] |
| 107 | """ |
| 108 | M, K = a.shape |
| 109 | assert a_compressed.shape == ( |
| 110 | M, |
| 111 | K // 2, |
| 112 | ), f"Expected a_compressed shape {(M, K // 2)}, got {a_compressed.shape}" |
| 113 | assert meta.shape == ( |
| 114 | M, |
| 115 | K // 4 // 8, |
| 116 | ), f"Expected meta shape {(M, K // 4 // 8)}, got {meta.shape}" |
| 117 | for m in range(M): |
| 118 | k_meta = 0 |
| 119 | for k in range(0, K, 4): |
| 120 | chunk = a[m, k : k + 4] |
| 121 | |
| 122 | non_zero_indices = torch.nonzero(chunk).squeeze() |
| 123 | meta_val = 0xE |
| 124 | if torch.equal(non_zero_indices, torch.tensor([0, 1])): |
| 125 | meta_val = 0x4 |
| 126 | elif torch.equal(non_zero_indices, torch.tensor([0, 2])): |
| 127 | meta_val = 0x8 |
| 128 | elif torch.equal(non_zero_indices, torch.tensor([0, 3])) or torch.equal( |
| 129 | non_zero_indices, torch.tensor(0) |
| 130 | ): |
| 131 | meta_val = 0xC |
| 132 | elif torch.equal(non_zero_indices, torch.tensor([1, 2])): |
| 133 | meta_val = 0x9 |
| 134 | elif torch.equal(non_zero_indices, torch.tensor([1, 3])) or torch.equal( |
| 135 | non_zero_indices, torch.tensor(1) |
| 136 | ): |
| 137 | meta_val = 0xD |
| 138 | elif torch.equal(non_zero_indices, torch.tensor([2, 3])) or torch.equal( |
| 139 | non_zero_indices, torch.tensor(2) |
| 140 | ): |
| 141 | meta_val = 0xE |
| 142 | elif torch.equal(non_zero_indices, torch.tensor([])) or torch.equal( |
| 143 | non_zero_indices, torch.tensor(3) |
| 144 | ): |
| 145 | meta_val = 0xE |
| 146 | else: |
| 147 | raise ValueError(f"Invalid non-zero pattern: {non_zero_indices}") |