import torch, gemlite
from torchao.quantization import GemliteUIntXWeightOnlyConfig, quantize_
device = 'cuda:0'
dtype = torch.float16
gemlite.set_autotune("default")

torch.manual_seed(0)
in_features, out_features, group_size = 256, 512, 64

orig_shape = [out_features, in_features]
layer = torch.nn.Linear(in_features, out_features, bias=False, dtype=dtype, device=device)
layer.weight.data /= 10.
weight = layer.weight.data.clone()

quantize_(layer, GemliteUIntXWeightOnlyConfig(bit_width=4, group_size=group_size))

meta_args =  layer.weight.tensor_impl.gemlite_kwargs['meta_args']
W_group_mode = meta_args[10]

#Test matmul
####################################################################################
torch.manual_seed(0)
x = torch.randn((1, layer.in_features), device=device, dtype=dtype) / 10.
y_ref = x @ weight.T
y_gem  = layer(x)
err = (y_ref - y_gem).abs().mean()
assert err < 5e-3, "Dot product mismatch. " + str(err)

#Test slicing 
####################################################################################
def dequant(input_layer, in_features, orig_shape):
    int_data = input_layer.tensor_impl.packed_weight
    scale = input_layer.tensor_impl.scale
    zero_point = input_layer.tensor_impl.zero_point

    W_q = (
        gemlite.bitpack.unpack_over_rows(
            int_data, W_nbits=4, num_output_rows=in_features, dtype=torch.uint8
        )
        .T.contiguous()
        .view([-1, group_size])
    )

    s = scale.t().contiguous().view(-1, 1)
    z = zero_point.t().contiguous().view(-1, 1)

    if W_group_mode == 4:  # FMA
        W_deq = (W_q * s + z).view(orig_shape)
    else:
        W_deq = ((W_q - z) * s).view(orig_shape)

    return W_deq

W_r = dequant(layer.weight, layer.in_features, orig_shape) #~weight


#Slicing in half
for slice_axis, start, end in [(0, 0, 256), (0, 256, 256), (1, 0, 128), (1, 128, 128)]:
    layer_sliced = layer.weight.narrow(slice_axis, start, end)

    if slice_axis == 0:
        num_rows, out_shape = layer.in_features, (orig_shape[0]//2, orig_shape[1]) 
    else:
        num_rows, out_shape = layer.in_features // 2, (orig_shape[0], orig_shape[1]//2)

    W_slice = dequant(layer_sliced, num_rows, out_shape)

    W_slice_ref = W_r[start:start+end, :] if slice_axis == 0 else W_r[:, start:start+end]
    assert (W_slice_ref - W_slice).abs().mean() == 0, f"Slicing {start}:{end} along axis={slice_axis} is incorrect"
    ```