| 1 | #ifndef CPU_MICRO_GEMM_INT8_VSX_HPP |
| 2 | #define CPU_MICRO_GEMM_INT8_VSX_HPP |
| 3 | |
| 4 | #include <algorithm> |
| 5 | #include <cstdint> |
| 6 | #include <cstring> |
| 7 | #include <type_traits> |
| 8 | #include <utility> |
| 9 | |
| 10 | #include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp" |
| 11 | |
| 12 | #include <altivec.h> |
| 13 | #include <c10/util/BFloat16.h> |
| 14 | #include <c10/util/Exception.h> |
| 15 | |
| 16 | namespace cpu_micro_gemm { |
| 17 | |
| 18 | namespace vsx_int8 { |
| 19 | |
| 20 | constexpr int32_t K = 4; |
| 21 | constexpr int32_t Mr = 4; |
| 22 | constexpr int32_t Nr = 32; |
| 23 | |
| 24 | // --------------------------------------------------------------------------- |
| 25 | // 8-Accumulator MMA 4x32 micro-GEMM (POWER10 xvi8ger4pp). |
| 26 | // |
| 27 | // Computes 4 rows x 32 columns (128 outputs) using all 8 hardware accumulators. |
| 28 | // Operates on pre-biased uint8 weights (B + 128) and subtracts 128 * a_sums. |
| 29 | // Works for all m in [1, 4] with zero scalar loop overhead. |
| 30 | // --------------------------------------------------------------------------- |
| 31 | __attribute__((target("cpu=power10"))) FORCE_INLINE void gemm_mma_4x32( |
| 32 | const int8_t* __restrict__ a_packed, const uint8_t* __restrict__ b_panel, |
| 33 | const int32_t* __restrict__ a_sums, int32_t* __restrict__ c_ptr, |
| 34 | const int32_t m, const int32_t k_size, const int64_t ldc) { |
| 35 | TORCH_CHECK(m > 0 && m <= Mr); |
| 36 | TORCH_CHECK(k_size % K == 0); |
| 37 | |
| 38 | __vector_quad acc0, acc1, acc2, acc3, acc4, acc5, acc6, acc7; |
| 39 | __builtin_mma_xxsetaccz(&acc0); |
| 40 | __builtin_mma_xxsetaccz(&acc1); |
| 41 | __builtin_mma_xxsetaccz(&acc2); |
| 42 | __builtin_mma_xxsetaccz(&acc3); |
| 43 | __builtin_mma_xxsetaccz(&acc4); |
| 44 | __builtin_mma_xxsetaccz(&acc5); |
| 45 | __builtin_mma_xxsetaccz(&acc6); |
| 46 | __builtin_mma_xxsetaccz(&acc7); |
| 47 | |
| 48 | const int8_t* __restrict__ ap = a_packed; |
| 49 | const uint8_t* __restrict__ bp = b_panel; |
| 50 | |
| 51 | #pragma GCC unroll 4 |
| 52 | for (int32_t k = 0; k < k_size; k += 4) { |
| 53 | const __vector unsigned char avec = |
| 54 | vec_xl(0, reinterpret_cast<const unsigned char*>(ap)); |
| 55 | __builtin_mma_xvi8ger4pp(&acc0, avec, vec_xl(0, bp)); |
| 56 | __builtin_mma_xvi8ger4pp(&acc1, avec, vec_xl(16, bp)); |
| 57 | __builtin_mma_xvi8ger4pp(&acc2, avec, vec_xl(32, bp)); |
| 58 | __builtin_mma_xvi8ger4pp(&acc3, avec, vec_xl(48, bp)); |
| 59 | __builtin_mma_xvi8ger4pp(&acc4, avec, vec_xl(64, bp)); |
| 60 | __builtin_mma_xvi8ger4pp(&acc5, avec, vec_xl(80, bp)); |
| 61 | __builtin_mma_xvi8ger4pp(&acc6, avec, vec_xl(96, bp)); |
| 62 | __builtin_mma_xvi8ger4pp(&acc7, avec, vec_xl(112, bp)); |
| 63 | ap += 16; |
| 64 | bp += 128; |
| 65 | } |
| 66 | |
| 67 | __vector signed int res[4]; |
| 68 | #define STORE_ACC(ACC, CG_IDX) \ |
| 69 | { \ |
| 70 | __builtin_mma_disassemble_acc(res, &ACC); \ |
| 71 | for (int r = 0; r < m; ++r) { \ |
| 72 | const __vector signed int vcorr = vec_splats(128 * a_sums[r]); \ |
| 73 | const __vector signed int val = vec_sub(res[r], vcorr); \ |
| 74 | vec_xst(val, 0, \ |
| 75 | reinterpret_cast<signed int*>(c_ptr + r * ldc + (CG_IDX) * 4)); \ |
| 76 | } \ |
| 77 | } |
| 78 | STORE_ACC(acc0, 0); |
| 79 | STORE_ACC(acc1, 1); |
| 80 | STORE_ACC(acc2, 2); |
| 81 | STORE_ACC(acc3, 3); |
| 82 | STORE_ACC(acc4, 4); |
| 83 | STORE_ACC(acc5, 5); |
| 84 | STORE_ACC(acc6, 6); |
| 85 | STORE_ACC(acc7, 7); |
| 86 | #undef STORE_ACC |
| 87 | } |
| 88 | |
| 89 | } // namespace vsx_int8 |
| 90 | |
| 91 | // ============================================================================= |
| 92 | // MicroGemmINT8 specialisation for VSX (POWER10+) |
| 93 | // ============================================================================= |
| 94 | template <typename scalar_t> |
| 95 | class MicroGemmINT8<cpu_utils::ISA::VSX, scalar_t> { |
| 96 | public: |
| 97 | static constexpr int32_t K = vsx_int8::K; // 4 |
| 98 | static constexpr int32_t Mr = vsx_int8::Mr; // 4 |
| 99 | static constexpr int32_t Nr = vsx_int8::Nr; // 32 |
| 100 | static constexpr int32_t NrGemv = Nr; |
| 101 | static constexpr int32_t MaxMSize = 8; |
| 102 | static constexpr int32_t NSize = Nr; |
| 103 | static constexpr int32_t WeightOCGroupSize = Nr; |
| 104 | static_assert(MaxMSize % Mr == 0); |
| 105 | |
| 106 | // ------------------------------------------------------------------------- |
| 107 | // quantize_row: scalar_t -> INT8 + per-token scale |
| 108 | // Single-pass 16-element vector conversion with aligned 16-byte stores. |
| 109 | // ------------------------------------------------------------------------- |
| 110 | static FORCE_INLINE void quantize_row(const scalar_t* input, int8_t* output, |
| 111 | float& scale, const int32_t size) { |
| 112 | TORCH_CHECK_EQ(size % 16, 0); |
| 113 | |
| 114 | const __vector signed short z = vec_splats((signed short)0); |
| 115 | |
| 116 | auto load_lo_hi = |
| 117 | [&](int32_t i) -> std::pair<__vector float, __vector float> { |
| 118 | if constexpr (std::is_same_v<scalar_t, float>) { |
| 119 | const float* p = reinterpret_cast<const float*>(input) + i; |
| 120 | return {vec_xl(0, p), vec_xl(16, p)}; |
| 121 | } else { |
| 122 | const signed short* p = |
| 123 | reinterpret_cast<const signed short*>(input + i); |
| 124 | const __vector signed short raw = (__vector signed short)vec_xl(0, p); |
| 125 | return {(__vector float)vec_mergeh(z, raw), |
| 126 | (__vector float)vec_mergel(z, raw)}; |
| 127 | } |
| 128 | }; |
| 129 | |
| 130 | __vector float max_v = vec_splats(0.0f); |
| 131 | for (int32_t i = 0; i < size; i += 8) { |
| 132 | auto [lo, hi] = load_lo_hi(i); |
| 133 | max_v = vec_max(max_v, vec_abs(lo)); |
| 134 | max_v = vec_max(max_v, vec_abs(hi)); |
| 135 | } |
| 136 | __vector float t = vec_max(max_v, vec_sld(max_v, max_v, 8)); |
| 137 | t = vec_max(t, vec_sld(t, t, 4)); |
| 138 | const float abs_max = std::max(vec_extract(t, 0), 1.0e-7f); |
| 139 | scale = abs_max / 127.0f; |
| 140 | const __vector float inv_scale = vec_splats(127.0f / abs_max); |
| 141 | |
| 142 | for (int32_t i = 0; i < size; i += 16) { |
| 143 | auto [lo0, hi0] = load_lo_hi(i); |
| 144 | auto [lo1, hi1] = load_lo_hi(i + 8); |
| 145 | |
| 146 | const __vector signed int i0 = |
| 147 | vec_cts(vec_round(vec_mul(lo0, inv_scale)), 0); |
| 148 | const __vector signed int i1 = |
| 149 | vec_cts(vec_round(vec_mul(hi0, inv_scale)), 0); |
| 150 | const __vector signed int i2 = |
| 151 | vec_cts(vec_round(vec_mul(lo1, inv_scale)), 0); |
| 152 | const __vector signed int i3 = |
| 153 | vec_cts(vec_round(vec_mul(hi1, inv_scale)), 0); |
| 154 | |
| 155 | const __vector signed short s0 = vec_packs(i0, i1); |
| 156 | const __vector signed short s1 = vec_packs(i2, i3); |
| 157 | const __vector signed char s8 = vec_packs(s0, s1); |
| 158 | |
| 159 | vec_xst(s8, 0, reinterpret_cast<signed char*>(output + i)); |
| 160 | } |
| 161 | } |
| 162 | |
| 163 | // ------------------------------------------------------------------------- |
| 164 | // dequantize_tile: INT32 accumulator -> FP32 scaled output |
| 165 | // ------------------------------------------------------------------------- |
| 166 | static FORCE_INLINE void dequantize_tile( |
| 167 | int32_t* input, float* output, const float* __restrict__ input_scales, |
| 168 | const float* __restrict__ weight_scales, const int32_t m, const int32_t n, |
| 169 | const int32_t stride) { |
| 170 | TORCH_CHECK_EQ(n % 4, 0); |
| 171 | for (int32_t mi = 0; mi < m; ++mi) { |
| 172 | const __vector float a_scale = vec_splats(input_scales[mi]); |
| 173 | for (int32_t ni = 0; ni < n; ni += 4) { |
| 174 | const __vector signed int i32 = vec_xl(0, input + mi * stride + ni); |
| 175 | const __vector float wscale = vec_xl(0, weight_scales + ni); |
| 176 | vec_xst(vec_mul(vec_ctf(i32, 0), vec_mul(a_scale, wscale)), 0, |
| 177 | output + mi * stride + ni); |
| 178 | } |
| 179 | } |
| 180 | } |
| 181 | |
| 182 | // ------------------------------------------------------------------------- |
| 183 | // pack_input_from_rows: pack rows into 4-row MMA tiles and compute row sums |
| 184 | // Header: 64 bytes (contains int32 a_sums[0..7]) |
| 185 | // Data: a_packed + 64 (contains 4x4 tiles of A) |
| 186 | // ------------------------------------------------------------------------- |
| 187 | static void pack_input_from_rows(const int8_t* const* __restrict__ rows, |
| 188 | int8_t* __restrict__ a_packed, |
| 189 | const int32_t m, const int32_t k) { |
| 190 | TORCH_CHECK(m > 0 && m <= MaxMSize); |
| 191 | TORCH_CHECK_EQ(k % K, 0); |
| 192 | |
| 193 | int32_t* a_sums = reinterpret_cast<int32_t*>(a_packed); |
| 194 | for (int32_t r = 0; r < 8; ++r) a_sums[r] = 0; |
| 195 | |
| 196 | int8_t* a_data = a_packed + 64; |
| 197 | |
| 198 | for (int32_t row_base = 0; row_base < m; row_base += Mr) { |
| 199 | const int32_t panel_m = std::min(Mr, m - row_base); |
| 200 | const int8_t* const* panel_rows = rows + row_base; |
| 201 | int32_t* panel_sums = a_sums + row_base; |
| 202 | int8_t* dst_panel = a_data + (row_base / Mr) * (k * Mr); |
| 203 | |
| 204 | for (int32_t ki = 0; ki < k; ki += 4) { |
| 205 | int8_t* dst = dst_panel + (ki / 4) * 16; |
| 206 | for (int32_t r = 0; r < 4; ++r) { |
| 207 | if (r < panel_m) { |
| 208 | const int8_t* src = panel_rows[r] + ki; |
| 209 | dst[r * 4 + 0] = src[0]; |
| 210 | dst[r * 4 + 1] = src[1]; |
| 211 | dst[r * 4 + 2] = src[2]; |
| 212 | dst[r * 4 + 3] = src[3]; |
| 213 | panel_sums[r] += |
| 214 | static_cast<int32_t>(src[0]) + src[1] + src[2] + src[3]; |
| 215 | } else { |
| 216 | dst[r * 4 + 0] = 0; |
| 217 | dst[r * 4 + 1] = 0; |
| 218 | dst[r * 4 + 2] = 0; |
| 219 | dst[r * 4 + 3] = 0; |
| 220 | } |
| 221 | } |
| 222 | } |
| 223 | } |
| 224 | } |
| 225 | |
| 226 | // ------------------------------------------------------------------------- |
| 227 | // pack_weight: [output_size, input_size] INT8 -> packed layout |
| 228 | // Pre-biases with +128 in [N/32][K/4][8 col-groups][4 cols x 4 k] layout |
| 229 | // ------------------------------------------------------------------------- |
| 230 | static void pack_weight(const int8_t* __restrict__ weight, |
| 231 | int8_t* __restrict__ packed_weight, |
| 232 | const int32_t output_size, const int32_t input_size) { |
| 233 | TORCH_CHECK_EQ(output_size % NSize, 0); |
| 234 | TORCH_CHECK_EQ(input_size % K, 0); |
| 235 | |
| 236 | uint8_t* dst_base = reinterpret_cast<uint8_t*>(packed_weight); |
| 237 | |
| 238 | for (int32_t o = 0; o < output_size; o += Nr) { |
| 239 | uint8_t* dst_panel = dst_base + o * input_size; |
| 240 | for (int32_t k = 0; k < input_size; k += 4) { |
| 241 | for (int32_t cg = 0; cg < 8; ++cg) { |
| 242 | uint8_t* dst = dst_panel + (k / 4) * 128 + cg * 16; |
| 243 | for (int32_t col = 0; col < 4; ++col) { |
| 244 | for (int32_t ki = 0; ki < 4; ++ki) { |
| 245 | const int8_t w = weight[(o + cg * 4 + col) * input_size + k + ki]; |
| 246 | dst[col * 4 + ki] = |
| 247 | static_cast<uint8_t>(static_cast<int>(w) + 128); |
| 248 | } |
| 249 | } |
| 250 | } |
| 251 | } |
| 252 | } |
| 253 | } |
| 254 | |
| 255 | // ------------------------------------------------------------------------- |
| 256 | // gemm: 8-accumulator MMA GEMM over 32-column tiles |
| 257 | // ------------------------------------------------------------------------- |
| 258 | void gemm(const int8_t* __restrict__ a_packed, |
| 259 | const int8_t* __restrict__ b_packed, int32_t* __restrict__ c, |
| 260 | const int32_t m, const int32_t k, const int64_t b_n_group_stride, |
| 261 | const int64_t ldc) const { |
| 262 | TORCH_CHECK(m > 0 && m <= MaxMSize); |
| 263 | TORCH_CHECK_EQ(k % K, 0); |
| 264 | (void)b_n_group_stride; |
| 265 | |
| 266 | const int32_t* a_sums = reinterpret_cast<const int32_t*>(a_packed); |
| 267 | const int8_t* a_data = a_packed + 64; |
| 268 | const uint8_t* b_u8 = reinterpret_cast<const uint8_t*>(b_packed); |
| 269 | |
| 270 | for (int32_t n_idx = 0; n_idx < NSize; n_idx += Nr) { |
| 271 | const uint8_t* b_panel = b_u8 + n_idx * k; |
| 272 | |
| 273 | for (int32_t row_base = 0; row_base < m; row_base += Mr) { |
| 274 | const int32_t panel_m = std::min(Mr, m - row_base); |
| 275 | const int8_t* a_panel = a_data + (row_base / Mr) * (k * Mr); |
| 276 | const int32_t* panel_sums = a_sums + row_base; |
| 277 | |
| 278 | vsx_int8::gemm_mma_4x32(a_panel, b_panel, panel_sums, |
| 279 | c + row_base * ldc + n_idx, panel_m, k, ldc); |
| 280 | } |
| 281 | } |
| 282 | } |
| 283 | }; |
| 284 | |
| 285 | } // namespace cpu_micro_gemm |
| 286 | |
| 287 | #endif // CPU_MICRO_GEMM_INT8_VSX_HPP |