| 680 | } |
| 681 | |
| 682 | Kernel selectMatmul(Context &ctx, int version, |
| 683 | const Bindings</* input, weights, output */ 3> &bindings, |
| 684 | size_t M, size_t K, size_t N, NumType numtype) { |
| 685 | Kernel kernel; |
| 686 | if (version == 1) { |
| 687 | Shape wgSize = {256, 1, 1}; |
| 688 | Shape nWorkgroups = cdiv({M, N, 1}, {16, 16, 1}); |
| 689 | KernelCode matmul = createNoOp(kShaderNoOp, /*wgsize*/ wgSize); |
| 690 | kernel = createKernel(ctx, matmul, bindings, |
| 691 | /*nWorkgroups*/ nWorkgroups); |
| 692 | } else if (version == 2) { |
| 693 | Shape wgSize = {16, 16, 1}; |
| 694 | LOG(kDefLog, kInfo, "wgSize: %s", toString(wgSize).c_str()); |
| 695 | KernelCode matmul = |
| 696 | createMatmul1(kShaderMatmul1, M, K, N, /*wgsize*/ wgSize, numtype); |
| 697 | kernel = createKernel(ctx, matmul, bindings, |
| 698 | /*nWorkgroups*/ cdiv({M, N, 1}, wgSize)); |
| 699 | } else if (version == 3) { |
| 700 | static constexpr size_t tileSize = 16; |
| 701 | KernelCode matmul = createMatmul2(kShaderMatmul2, M, K, N, |
| 702 | /*wgSize*/ {tileSize * tileSize, 1, 1}, numtype); |
| 703 | kernel = |
| 704 | createKernel(ctx, matmul, bindings, |
| 705 | /* nWorkgroups*/ cdiv({M, N, 1}, {tileSize, tileSize, 1})); |
| 706 | } else if (version == 4 || version == 6) { |
| 707 | static constexpr size_t BM = 64; |
| 708 | static constexpr size_t BK = 4; |
| 709 | static constexpr size_t BN = BM; |
| 710 | static constexpr size_t TM = |
| 711 | BN / BK; // BM * BN / TM == BM * BK, therefore TM == BN / BK |
| 712 | Shape wgSize = {BM * BN / TM, 1, |
| 713 | 1}; // BM * BN values per workgroup, TM values per thread |
| 714 | Shape nWorkgroups = {cdiv(M, BM), cdiv(N, BN), 1}; |
| 715 | LOG(kDefLog, kInfo, "M: %d, K: %d, N: %d", M, K, N); |
| 716 | LOG(kDefLog, kInfo, "BM: %d, BK: %d, BN: %d, TM: %d", BM, BK, BN, TM); |
| 717 | LOG(kDefLog, kInfo, "wgSize: ( %s )", toString(wgSize).c_str()); |
| 718 | LOG(kDefLog, kInfo, "nWorkgroups: ( %s )", toString(nWorkgroups).c_str()); |
| 719 | KernelCode matmul = createMatmul3(kShaderMatmul3, M, K, N, BM, BK, BN, TM, |
| 720 | /*wgSize*/ wgSize, |
| 721 | numtype, |
| 722 | /*Loop unrolling*/ version == 6 ? true: false); |
| 723 | kernel = createKernel(ctx, matmul, bindings, |
| 724 | /*nWorkgroups*/ nWorkgroups); |
| 725 | } else if (version == 5 || version == 7) { |
| 726 | static constexpr size_t BM = 64; |
| 727 | static constexpr size_t BK = 8; |
| 728 | static constexpr size_t BN = 64; |
| 729 | static constexpr size_t TM = BM / BK; |
| 730 | static constexpr size_t TN = BN / BK; |
| 731 | Shape wgSize = {(BM / TM) * (BN / TN), 1, 1}; // This is the same as BK * BK. |
| 732 | Shape nWorkgroups = {cdiv(M, BM), cdiv(N, BN), 1}; |
| 733 | LOG(kDefLog, kInfo, "M: %d, K: %d, N: %d", M, K, N); |
| 734 | LOG(kDefLog, kInfo, "BM: %d, BK: %d, BN: %d, TM: %d, TN: %d", BM, BK, BN, TM, TN); |
| 735 | LOG(kDefLog, kInfo, "wgSize: ( %s )", toString(wgSize).c_str()); |
| 736 | LOG(kDefLog, kInfo, "nWorkgroups: ( %s )", toString(nWorkgroups).c_str()); |
| 737 | KernelCode matmul = createMatmul4(kShaderMatmul4, M, K, N, BM, BK, BN, TM, TN, |
| 738 | /*wgSize*/ wgSize, |
| 739 | numtype, |
no test coverage detected