| 721 | return fp32_arr.view(np.float32) |
| 722 | |
| 723 | def preprocess_weights( |
| 724 | w: np.ndarray, |
| 725 | bits = 2, |
| 726 | g = 4, |
| 727 | ) -> Tuple[np.ndarray, np.ndarray]: |
| 728 | M, K = w.shape |
| 729 | |
| 730 | cf=configparser.ConfigParser() |
| 731 | cf.read("./build/kcfg.ini") |
| 732 | secs=cf.sections() |
| 733 | for sec in secs: |
| 734 | sec_splits = str(sec).split('_') |
| 735 | if sec_splits[-4] == "m" + str(M*2) and sec_splits[-3] == "k" + str(K): |
| 736 | bm = int(cf.get(sec, 'bm')) |
| 737 | kfactor = int(cf.get(sec, 'kfactor')) |
| 738 | simd_n_in = int(cf.get(sec, 'simd_n_in')) |
| 739 | simd_n_out = int(cf.get(sec, 'simd_n_out')) |
| 740 | break |
| 741 | |
| 742 | M = M * bits |
| 743 | ngroups_per_elem = 8 // g |
| 744 | |
| 745 | # (M // bits, K, bits) |
| 746 | w = np.stack([(w >> ib) & 1 for ib in range(bits)], axis=-1) |
| 747 | # print(w) |
| 748 | # (M // bits, K, bits) -> (M // bits, bits, K) -> (M // bits, bits, K) -> (M // bits, bits, K // g, g) |
| 749 | w = w.transpose(0, 2, 1).reshape(M // bits, bits, K // g, g) |
| 750 | w = sum([(w[:, :, :, ig] << ig) for ig in range(g)]) |
| 751 | # print(w) |
| 752 | # 0, 16, 1, 17, 2, 18, 3, 19, 4, 20, 5, 21, 6, 22, 7, 23, 8, 24, 9, 25, 10, 26, 11, 27, 12, 28, 13, 29, 14, 30, 15, 31 |
| 753 | # for bits=3 |
| 754 | # bit0: [0, 8), bit1: [8, 16), bit2: [16, 24), bit0: [24, 32) |
| 755 | # (M // bits // simd_n_float16, bits, simd_n_float16, K // g) |
| 756 | w = w.reshape(M // bits // simd_n_out, simd_n_out, bits, K // g).transpose(0, 2, 1, 3) |
| 757 | mgroup = ngroups_per_elem * simd_n_in |
| 758 | w = w.reshape(M // mgroup, ngroups_per_elem, simd_n_in, K // g).transpose(0, 2, 1, 3) |
| 759 | # 0 1 2 3 4 5 |
| 760 | w = w.reshape(M // bm, bm // mgroup, simd_n_in, ngroups_per_elem, K // g // kfactor, kfactor).transpose(0, 4, 1, 5, 2, 3) |
| 761 | w = sum([(w[:, :, :, :, :, ng] << (ng * g)) for ng in range(ngroups_per_elem)]) |
| 762 | w = w.reshape(M // bm, K // g // kfactor, bm // mgroup, kfactor, simd_n_in) |
| 763 | # input size of current TVM API |
| 764 | w = w.reshape(M // bm, K // g, bm // ngroups_per_elem) |
| 765 | |
| 766 | return w |
| 767 | |
| 768 | def transform_to_i2(x : NDArray): |
| 769 | x_num = np.prod(x.shape) |