Function preprocess_two_weights_tl2

utils/generate-dummy-bitnet-model.py:577–618

(M, K, weight_num, BM, BY, bm, by, weight, final_weight)

def preprocess_two_weights_tl2(M, K, weight_num, BM, BY, bm, by, weight, final_weight):
    weight = np.reshape(weight, (weight_num // 2, 2))
    hi_weight = np.multiply(np.split(weight, 2, axis=1)[0], 3)
    lo_weight = np.split(weight, 2, axis=1)[1]

    weight = np.reshape((hi_weight + lo_weight), weight_num // 2)

    # row-major index
    weight = weight + 4
    weight = np.reshape(weight, (M, K // 2)).astype(np.uint8)

    outer_BM_weights = np.split(weight, (M // BM), axis=0)
    for outer_BM_weight in outer_BM_weights:
        # split in col with size of by (32index * 3 == 96nums)
        outer_BY_weights = np.split(outer_BM_weight, (K // BY), axis=1)
        for outer_BY_weight in outer_BY_weights:
            # split in row with size of bm (32)
            inner_bm_weights = np.split(outer_BY_weight, (BM // bm), axis=0)
            for inner_bm_weight in inner_bm_weights:
                # split in col with size of by (2index * 2 == 4nums)
                inner_by_weights = np.split(inner_bm_weight, (BY // by), axis=1)
                for inner_by_weight in inner_by_weights:
                    func_weights = np.split(inner_by_weight, 2, axis=1)

                    left_weight = func_weights[0]
                    left_sub_weights = np.split(left_weight, 4, axis=0)
                    new_left_weight = np.reshape(
                                        np.concatenate([left_sub_weights[0], left_sub_weights[2], 
                                        left_sub_weights[1], left_sub_weights[3]], axis=0, dtype=np.uint8),
                                        (bm))

                    right_weight = func_weights[1]
                    right_sub_weights = np.split(right_weight, 4, axis=0)
                    new_right_weight = np.reshape(
                                        np.concatenate([right_sub_weights[0], right_sub_weights[2], 
                                        right_sub_weights[1], right_sub_weights[3]], axis=0, dtype=np.uint8),
                                        (bm))
                    hi_weight = new_left_weight.astype(np.uint8) << 4
                    lo_weight = new_right_weight
                    func_weight = hi_weight + lo_weight
                    func_weight = np.reshape(func_weight, bm * by // 4)
                    final_weight.append(func_weight)

def preprocess_three_weights_tl2(M, K, weight_num, BM, BY, bm, by, weight, final_weight):
    weight = np.reshape(weight, (weight_num // 3, 3))