| 530 | return kernel_code |
| 531 | |
| 532 | def gen_top_api(kernel_shapes, k_list): |
| 533 | |
| 534 | kernel_code = "void ggml_preprocessor(int bs, int m, int three_k, int two_k, void* B, void* LUT_Scales, void* Three_QLUT, void* Two_QLUT) {{\n\ |
| 535 | partial_max_reset(bs, (&(((float*)LUT_Scales)[0])));\n\ |
| 536 | if (m == {0} && two_k == {1} && three_k == {2}) {{\n\ |
| 537 | for (int32_t b = 0; b < bs; b++) {{\n\ |
| 538 | per_tensor_quant(two_k + three_k, (&(((float*)LUT_Scales)[b])), (&(((float*)B)[b * (two_k + three_k)])));\n\ |
| 539 | three_lut_ctor<{2}>((&(((int8_t*)Three_QLUT)[b * three_k / 3 * 32])), (&(((float*)B)[b * (three_k + two_k)])), (&(((float*)LUT_Scales)[b])));\n\ |
| 540 | two_lut_ctor<{1}>((&(((int8_t*)Two_QLUT)[b * two_k / 2 * 32])), (&(((float*)B)[b * (three_k + two_k) + {2}])), (&(((float*)LUT_Scales)[b])));\n\ |
| 541 | }}\n\ |
| 542 | }}\n\ |
| 543 | ".format(kernel_shapes[0][0], k_list[0][0], k_list[0][1]) |
| 544 | for i in range(1, len(kernel_shapes)): |
| 545 | kernel_code = "".join([kernel_code, " else if (m == {0} && two_k == {1} && three_k == {2}) {{\n\ |
| 546 | for (int32_t b = 0; b < bs; b++) {{\n\ |
| 547 | per_tensor_quant(two_k + three_k, (&(((float*)LUT_Scales)[b])), (&(((float*)B)[b * (two_k + three_k)])));\n\ |
| 548 | three_lut_ctor<{2}>((&(((int8_t*)Three_QLUT)[b * three_k / 3 * 32])), (&(((float*)B)[b * (three_k + two_k)])), (&(((float*)LUT_Scales)[b])));\n\ |
| 549 | two_lut_ctor<{1}>((&(((int8_t*)Two_QLUT)[b * two_k / 2 * 32])), (&(((float*)B)[b * (three_k + two_k) + {2}])), (&(((float*)LUT_Scales)[b])));\n\ |
| 550 | }}\n\ |
| 551 | }}\n".format(kernel_shapes[i][0], k_list[i][0], k_list[i][1])]) |
| 552 | kernel_code = "".join([kernel_code, "}\n"]) |
| 553 | |
| 554 | |
| 555 | kernel_code = "".join([kernel_code, "void ggml_qgemm_lut(int bs, int m, int k, int BK, void* A, void* sign, void* LUT, void* Scales, void* LUT_Scales, void* C) {{\n\ |
| 556 | if (m == {0} && k == {1}) {{\n\ |
| 557 | if (BK == {2}) {{\n\ |
| 558 | if (bs == 1) {{\n\ |
| 559 | two_qgemm_lut_{4}<1>(A, LUT, Scales, LUT_Scales, C);\n\ |
| 560 | }} else if (bs == 8) {{\n\ |
| 561 | two_qgemm_lut_{4}<8>(A, LUT, Scales, LUT_Scales, C);\n\ |
| 562 | }} else if (bs == 32) {{\n\ |
| 563 | two_qgemm_lut_{4}<32>(A, LUT, Scales, LUT_Scales, C);\n\ |
| 564 | }} else if (bs == 128) {{\n\ |
| 565 | two_qgemm_lut_{4}<128>(A, LUT, Scales, LUT_Scales, C);\n\ |
| 566 | }} else if (bs == 256) {{\n\ |
| 567 | two_qgemm_lut_{4}<256>(A, LUT, Scales, LUT_Scales, C);\n\ |
| 568 | }} else if (bs == 512) {{\n\ |
| 569 | two_qgemm_lut_{4}<512>(A, LUT, Scales, LUT_Scales, C);\n\ |
| 570 | }}\n\ |
| 571 | }}\n\ |
| 572 | else if (BK == {3}) {{\n\ |
| 573 | if (bs == 1) {{\n\ |
| 574 | three_qgemm_lut_{4}<1>(A, sign, LUT, Scales, LUT_Scales, C);\n\ |
| 575 | }}else if (bs == 8) {{\n\ |
| 576 | three_qgemm_lut_{4}<8>(A, sign, LUT, Scales, LUT_Scales, C);\n\ |
| 577 | }}else if (bs == 32) {{\n\ |
| 578 | three_qgemm_lut_{4}<32>(A, sign, LUT, Scales, LUT_Scales, C);\n\ |
| 579 | }}else if (bs == 128) {{\n\ |
| 580 | three_qgemm_lut_{4}<128>(A, sign, LUT, Scales, LUT_Scales, C);\n\ |
| 581 | }}else if (bs == 256) {{\n\ |
| 582 | three_qgemm_lut_{4}<256>(A, sign, LUT, Scales, LUT_Scales, C);\n\ |
| 583 | }}else if (bs == 512) {{\n\ |
| 584 | three_qgemm_lut_{4}<512>(A, sign, LUT, Scales, LUT_Scales, C);\n\ |
| 585 | }}\n\ |
| 586 | }}\n\ |
| 587 | }}\n\ |
| 588 | ".format(kernel_shapes[0][0], kernel_shapes[0][1], k_list[0][0], k_list[0][1], "{}_{}".format(kernel_shapes[0][0], kernel_shapes[0][1]))]) |
| 589 | for i in range(1, len(kernel_shapes)): |