(pre, BM, BK, bm, k)
| 222 | return all_code |
| 223 | |
| 224 | def gen_tbl_impl(pre, BM, BK, bm, k): |
| 225 | |
| 226 | kernel_code = "\ |
| 227 | #include <arm_neon.h>\n\ |
| 228 | \n\ |
| 229 | #define BM{0} {1}\n\ |
| 230 | #define BBK{0} {2}\n\ |
| 231 | inline void tbl_impl_{0}(int32_t* c, int8_t* lut, uint8_t* a) {{\n\ |
| 232 | #ifdef __ARM_NEON\n\ |
| 233 | const int KK = BBK{0} / 2;\n\ |
| 234 | const uint8x16_t vec_mask = vdupq_n_u8(0x0f);\n\ |
| 235 | const int8x16_t vec_zero = vdupq_n_s16(0x0000);\n\ |
| 236 | int8x16_t vec_lut[2 * KK];\n\ |
| 237 | ".format(pre, BM, BK) |
| 238 | |
| 239 | kernel_code = "".join([kernel_code, " int16x8_t vec_c[{}];".format(bm // 8)]) |
| 240 | |
| 241 | kernel_code = "".join([kernel_code, "\n\ |
| 242 | #pragma unroll\n\ |
| 243 | for (int k = 0; k < 2 * KK; k++) {\n\ |
| 244 | vec_lut[k] = vld1q_s8(lut + k * 16);\n\ |
| 245 | }\n"]) |
| 246 | |
| 247 | pre_core_code = "\n\ |
| 248 | #pragma unroll\n\ |
| 249 | for (int i = 0; i < BM{}; i += {}) {{\n\ |
| 250 | #pragma unroll\n\ |
| 251 | for (int i=0; i<{}; i++) {{\n\ |
| 252 | vec_c[i] = vandq_s16(vec_c[i], vec_zero);\n\ |
| 253 | }}\n".format(pre, bm, bm // 8) |
| 254 | |
| 255 | body_core_pre_code = "\n\ |
| 256 | #pragma unroll\n\ |
| 257 | for (int k = 0; k < KK / {}; k++) {{\n\ |
| 258 | ".format(256 // bm // 2) |
| 259 | |
| 260 | body_core_post_code = "\n\ |
| 261 | }\n\ |
| 262 | \ |
| 263 | #endif\n\ |
| 264 | }\n" |
| 265 | |
| 266 | kernel_code = "".join([kernel_code, pre_core_code, body_core_pre_code, gen_body_core_code(bm, 256 // bm), body_core_post_code]) |
| 267 | |
| 268 | kernel_code = "".join([kernel_code, "\n\ |
| 269 | int32_t qgemm_lut_{0}(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {{\n\ |
| 270 | alignas({1}) uint32_t CBits[BM{0}];\n\ |
| 271 | memset(&(CBits[0]), 0, BM{0} * sizeof(int32_t));\n\ |
| 272 | #pragma unroll\n\ |
| 273 | for (int32_t k_outer = 0; k_outer < {2} / BBK{0}; ++k_outer) {{\n\ |
| 274 | tbl_impl_{0}((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK{0} / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK{0} / 2 / 2 * BM{0})])));\n\ |
| 275 | }}\n\ |
| 276 | #pragma unroll\n\ |
| 277 | for (int i = 0; i < BM{0}; i++) {{\n\ |
| 278 | ((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];\n\ |
| 279 | }}\n\ |
| 280 | return 0;\n\ |
| 281 | }};\n".format(pre, min(32, BK), k)]) |
no test coverage detected