group_size = 128
input_shape = (1, 1024)
weight_shape = (1024, 1024)
scaling_shape = (1024, 1024 // 128)
zeros_shape = (1024, 1024 // 128)
output_shape = (1, 1024)

scaling = torch.rand(scaling_shape, dtype=torch.float16).cuda()
zeros = torch.rand(zeros_shape, dtype=torch.float16).cuda()

# Compute reference result with manual scaling and zero-point adjustment
# rescale = (weight - zeros) * scaling
for i in range(in_features // group_size): # group number i in range(8)
    for j in range(group_size): # group j-th element/column
         # within each group, we use the same zeros and scaling factors.
         rescaling_tensor[:, i*group_size+j] = (weight_tensor[:, i*group_size+j] - zeros[:, i]) * scaling[:, i]

# from tvm.script import ir as I
# from tvm.script import tir as T

@I.ir_module
class Module:
    @T.prim_func
    def main(
      A: T.Buffer((1, 1024), "float16"),
      B: T.Buffer((1024, 512), "int8"), Scale: T.Buffer((1024, 8), "float16"),
      Zeros: T.Buffer((1024, 8), "float16"),
      D: T.Buffer((1, 1024), "float16")):
        # with T.block("root"):
        B_decode = T.alloc_buffer((1024, 1024), "float16")
        C = T.alloc_buffer((1, 1024), "float16")
        for n, k in T.grid(1024, 1024):
            with T.block("B_decode"):
                v_n, v_k = T.axis.remap("SS", [n, k]) # “S” (for spatial), “R” (for reduction)
                T.reads(B[v_n, v_k // 2], Zeros[v_n, v_k // 128], Scale[v_n, v_k // 128])
                T.writes(B_decode[v_n, v_k])
                B_decode[v_n, v_k] = # decompressing B
                  (
                     T.Cast("float16", T.bitwise_and(
                        T.shift_right(B[v_n, v_k // 2], T.Cast("int8", v_k % 2 * 4)),
                        T.int8(15) # b1111
                     ))
                     -
                     Zeros[v_n, v_k // 128] # re-centering
                  )
                  * Scale[v_n, v_k // 128] # scaling 
        for i, j, k in T.grid(1, 1024, 1024):
            with T.block("C"):
                v_i, v_j, v_k = T.axis.remap("SSR", [i, j, k])
                T.reads(A[v_i, v_k], B_decode[v_j, v_k])
                T.writes(C[v_i, v_j])
                with T.init():
                    C[v_i, v_j] = T.float16(0)
                C[v_i, v_j] = C[v_i, v_j] + A[v_i, v_k] * B_decode[v_j, v_k] # matrix multiplication
        for i, j in T.grid(1, 1024):
            with T.block("D"):
                v_i, v_j = T.axis.remap("SS", [i, j])
                T.reads(C[v_i, v_j])
                T.writes(D[v_i, v_j])
                D[v_i, v_j] = C[v_i, v_j]

import sys
import os
path = os.path.join("/home/tk/Desktop/bitblas/BitBLAS", "./build/lib")
sys.path.insert(0, path)
import torch
import bitblas

class TvmLinear(torch.nn.Module):
    def __init__(self, in_features, out_features, W_dtype="uint4"):
        matmul_config = bitblas.MatmulConfig(
            M=1,
            K=in_features,
            N=out_features,
            A_dtype="float16",  # activation A dtype
            W_dtype=W_dtype,  # weight W dtype
            accum_dtype="float16",  # accumulation dtype
            out_dtype="float16",  # output dtype
            layout="nt", # A is non-transpose and W is transpose
            with_bias=False,
            group_size=-1, # 128,  # setting for grouped quantization
            with_scaling=False,  # setting for scaling factor
            with_zeros=False,  # setting for zeros
            zeros_mode="original",  # setting for how to calculating zeros
        )
        self.matmul = bitblas.Matmul(config=matmul_config)
        init_W = torch.randint(0, 7, (out_features, in_features), dtype=torch.int8).cuda()
        self.set_weight(init_W)

    def set_weight(self, origin_int_W):
        self.W = self.matmul.transform_weight(origin_int_W)
        self.W_ori = origin_int_W

    def forward(self, A):
        output = self.matmul(A, self.W)
        verify = A @ self.W_ori.T.half()
        assert torch.allclose(output, verify, atol=1e-2)
        return output


inp = torch.rand((1, 8), dtype=torch.float16).cuda()
new_module = TvmLinear(8, 5)
out = new_module.forward(inp)

class TvmLinear(torch.nn.Module):
    def __init__(self, batch_size, in_features, out_features,
        W_dtype="uint4", group_size=-1, debug=False, no_extra_mem=False, tuning=True):
        super().__init__() # set up torch module (e.g., _backward_hooks)

        matmul_config = bitblas.MatmulConfig(
            M=batch_size,
            K=in_features,
            N=out_features,
            A_dtype="float16",  # activation A dtype
            W_dtype=W_dtype,  # weight W dtype
            storage_dtype='int8',
            accum_dtype="float16",  # accumulation dtype
            out_dtype="float16",  # output dtype
            layout="nt",
            with_bias=False,
            group_size=group_size,
            with_scaling=True,  # setting for scaling factor
            with_zeros=True,  # setting for zeros
            zeros_mode="rescale",  # setting for how to calculating zeros
            fast_decoding=False # important! avoid post-processing (i.e., LOP3Permutate)
        )
        self.group_size = group_size if group_size != -1 else in_features
        self.matmul = bitblas.Matmul(config=matmul_config, enable_tuning=tuning)

        # set random initial (binary) weights
        #init_W = torch.randint(0, 2, (out_features, in_features), dtype=torch.int8)
        #self._set_quantized_weight(init_W)

        self.bits = int(''.join(list(filter(str.isdigit, W_dtype))))
        self.scaling = None
        self.zeros = None
        self.weight = None # placeholder for T5 modeling access.
        self.debug = debug
        self.no_extra_mem = no_extra_mem

    def _set_quantized_weight(self, W_quant):
        self.W_store = self.matmul.transform_weight(W_quant)
        self.W_store = self.W_store.cuda()

    def set_weight(self, W):
        out_features, in_features = W.shape
        assert in_features % self.group_size == 0

        reshape = (out_features, -1, self.group_size)
        W_reshape = W.reshape(*reshape)
        group_max = W_reshape.max(-1).values
        group_min = W_reshape.min(-1).values
        group_max = group_max.unsqueeze(-1).expand(reshape)
        group_min = group_min.unsqueeze(-1).expand(reshape)
        #Q_min = -(2**(self.bits - 1))
        #Q_max = 2**(self.bits - 1) - 1
        Q_min = 0
        Q_max = 2**(self.bits) - 1
        ratio = (Q_max - Q_min) / (group_max - group_min)

        W_quant = ((W_reshape - group_min) * ratio).round() + Q_min
        W_quant = W_quant.to(dtype=torch.long)
        W_quant = W_quant.reshape(W.shape)

        self._set_quantized_weight(W_quant)
        if not self.no_extra_mem:
            self.scaling = 1 / ratio[:,:,0].clone().cuda()
            self.zeros = Q_min * self.scaling - group_min[:,:,0].clone().cuda()

        if self.debug:
            self.W_origin = W
            self.W_quant = W_quant
            self.debug_scaling = self.scaling.unsqueeze(-1).expand(reshape).reshape(W.shape)
            self.debug_zeros = self.zeros.unsqueeze(-1).expand(reshape).reshape(W.shape)
            self.debug_W = W_quant.float() * self.debug_scaling - self.debug_zeros

    def forward(self, A):
        output = self.matmul(A, self.W_store, scale=self.scaling, zeros=self.zeros)
        if self.debug:
            output_check = A @ self.W_origin.T
            print(output.isnan().nonzero().numel())
            print(torch.allclose(self.debug_W.half(), self.W_origin, atol=1e-1))
            print(torch.allclose(A @ self.debug_W.half().T, A @ self.W_origin.T, atol=1e-1))
            breakpoint()
        return output

inp = torch.rand((4, 1, 8), dtype=torch.float16).cuda()
M = TvmLinear(4, 8, 5, W_dtype="uint4", group_size=4)

W = torch.rand((5, 8), dtype=torch.float16).cuda()
M.set_weight(W)

out = M.forward(inp)
print('inp', inp)
print('W', W)
print('out', out)

assert torch.allclose(
    M.debug_W.T.half(),
    W.T,
    atol=1e-1
)
assert torch.allclose(
    inp @ M.debug_W.T.half(),
    inp @ W.T,
    atol=1e-1
)
assert torch.allclose(
    inp @ W.T,
    out,
    atol=1e-1
)

from transformers import AutoProcessor, MusicgenForConditionalGeneration
model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small",
    torch_dtype=torch.half)
model.to("cuda")
print(model)

processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
inputs = processor(
    text=["80s pop track with bassy drums and synth"],
    padding=True, return_tensors="pt",
)
inputs.to("cuda")
seq_len = inputs.input_ids.shape[-1]

linear_classes = (torch.nn.Linear, )
to_replace_lst = []
for key, module in model.named_modules():
    if isinstance(module, linear_classes):
        if module.bias is not None: continue
        old_W = module.weight.detach().cpu()
        path_fields = key.split('.')
        parent_key = '.'.join(path_fields[:-1])
        child_key = path_fields[-1]
        to_replace_lst.append((parent_key, child_key, old_W))

n_replaced = 0
for parent_key, child_key, old_W in to_replace_lst:
    parent = model.get_submodule(parent_key)
    delattr(parent, child_key)
    out_features, in_features = old_W.shape
    M = TvmLinear([2, seq_len], in_features, out_features,
        group_size=-1, W_dtype="uint4", no_extra_mem=False, tuning=False)
    # 589824 * 2 => 294912 * 1 + 768 * 2 + 768 * 2
    # save: 881,664 bytes
    M.set_weight(old_W)
    setattr(parent, child_key, M)
    n_replaced += 1
    print(n_replaced, parent_key, child_key, old_W.shape)
    if n_replaced % 20 == 0:
        print('GC', gc.collect())
        torch.cuda.empty_cache()

breakpoint()

print('generating ...')
with torch.no_grad():
    audio_values = model.generate(**inputs, max_new_tokens=256)

import scipy
sampling_rate = model.config.audio_encoder.sampling_rate
data = audio_values[0, 0].detach().cpu().float().numpy()
scipy.io.wavfile.write("output.wav", rate=sampling_rate, data=data)

import sys
import os
path = os.path.join("/home/tk/Desktop/bitblas/BitBLAS", "./build/lib")
path2 = os.path.join("/home/tk/Desktop/bitblas/BitBLAS", "./build/lib/bitblas/3rdparty/tvm/python")
sys.path.insert(0, path)
sys.path.insert(0, path2)

import tvm
from tvm.ir.module import IRModule
from tvm.script import tir as T
import numpy as np

@T.prim_func
def main(var_A: T.handle, B: T.Buffer((768, 384), "int8"), Scale: T.Buffer((768, 3), "float16"), Zeros: T.Buffer((768, 3), "float16"), var_D: T.handle):
    T.func_attr({"dequantize_info": {"B_decode": {"decode_block": "B_decode", "fast_decoding": T.bool(False), "group_size": 256, "source_format": {"bits": 4, "format": "uint"}, "storage_dtype": "int8", "target_format": "float16", "with_scaling": T.bool(True), "with_zeros": T.bool(True), "zeros_mode": "rescale"}}, "opt_shapes": {"m": [2, 12]}, "tir.noalias": T.bool(True)})
    m = T.int32()
    A = T.match_buffer(var_A, (m, 768), "float16")
    D = T.match_buffer(var_D, (m, 768), "float16")
    # with T.block("root"):
    B_decode = T.alloc_buffer((768, 768), "float16")
    C = T.alloc_buffer((m, 768), "float16")
    for n, k in T.grid(768, 768):
        with T.block("B_decode"):
            v_n, v_k = T.axis.remap("SS", [n, k])
            T.reads(B[v_n, v_k // 2], Scale[v_n, v_k // 256], Zeros[v_n, v_k // 256])
            T.writes(B_decode[v_n, v_k])
            B_decode[v_n, v_k] = T.Cast("float16", T.bitwise_and(T.shift_right(B[v_n, v_k // 2], T.Cast("int8", v_k % 2 * 4)), T.int8(15))) * Scale[v_n, v_k // 256] - Zeros[v_n, v_k // 256]
    for i, j, k in T.grid(m, 768, 768):
        with T.block("C"):
            v_i, v_j, v_k = T.axis.remap("SSR", [i, j, k])
            T.reads(A[v_i, v_k], B_decode[v_j, v_k])
            T.writes(C[v_i, v_j])
            with T.init():
                C[v_i, v_j] = T.float16(0)
            C[v_i, v_j] = C[v_i, v_j] + A[v_i, v_k] * B_decode[v_j, v_k]
    for i, j in T.grid(m, 768):
        with T.block("D"):
            v_i, v_j = T.axis.remap("SS", [i, j])
            T.reads(C[v_i, v_j])
            T.writes(D[v_i, v_j])
            D[v_i, v_j] = C[v_i, v_j]

rt_mod = tvm.build(main, target='llvm')
func = rt_mod[rt_mod.entry_name]

A = tvm.nd.array(np.ones((2, 768), dtype="float16"))
B = tvm.nd.array(np.random.randint(0, 2, size=(768, 384), dtype="int8"))
Scale = tvm.nd.array(np.ones((768, 3), dtype="float16"))
Zeros = tvm.nd.array(np.ones((768, 3), dtype="float16"))
D = tvm.nd.array(np.ones((2, 768), dtype="float16"))
func(A, B, Scale, Zeros, D)
print(D)