import torch, cutlass, cutlass.cute as cute, cutlass.utils as utils
from cutlass.cute.nvgpu import tcgen05
from cutlass import Float32, BFloat16
from cutlass.utils import LayoutEnum
from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset
import cuda.bindings.driver as cuda
import cutlass.torch as ct

m, n, head_dim = 128, 128, 64
q = torch.randn(m, head_dim, 1, dtype=torch.bfloat16, device='cuda')
k = torch.randn(n, head_dim, 1, dtype=torch.bfloat16, device='cuda')
v_base = torch.randn(head_dim, n, dtype=torch.bfloat16, device='cuda')
v = v_base.as_strided((head_dim, n), (1, head_dim)).unsqueeze(-1)
c = torch.zeros(m, head_dim, 1, dtype=torch.bfloat16, device='cuda')

mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q))
mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k))
mV = ct.from_dlpack(v).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v))
mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c))
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)

class InspectKernel:
    def __init__(self):
        self.q_dtype = BFloat16; self.acc_dtype = Float32

    @cute.jit
    def __call__(self, q, k, v, c, stream):
        a_major = LayoutEnum.from_tensor(q).mma_major_mode()
        b_major = LayoutEnum.from_tensor(k).mma_major_mode()
        v_major = LayoutEnum.from_tensor(v).mma_major_mode()

        qk_mma = utils.sm100.make_trivial_tiled_mma(
            self.q_dtype, self.q_dtype, a_major, b_major,
            self.acc_dtype, tcgen05.CtaGroup.ONE, (128, 128), tcgen05.OperandSource.SMEM)
        pv_mma = utils.sm100.make_trivial_tiled_mma(
            self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, v_major,
            self.acc_dtype, tcgen05.CtaGroup.ONE, (128, 64), tcgen05.OperandSource.TMEM)

        qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2])
        qk_mma_tiler = (128, 128, qk_inst_k * 4)
        pv_mma_tiler = (qk_mma_tiler[0], qk_mma_tiler[2], qk_mma_tiler[1])

        qk_thr = qk_mma.get_slice(0)
        pv_thr = pv_mma.get_slice(0)

        qk_acc_shape = qk_thr.partition_shape_C(qk_mma_tiler[:2])
        tStS = qk_thr.make_fragment_C(qk_acc_shape)

        p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, pv_mma_tiler, self.q_dtype, 1)
        tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer)
        tOrP_base = pv_thr.make_fragment_A(tP)
        tOrP = tOrP_base[(None, None, None, 0)]

        tmem_p0_offset = 32
        p_offset = self.acc_dtype.width // self.q_dtype.width * tmem_p0_offset

        # Print pointer values - if tOrP inherits FP32 type, +64 adds 256 bytes
        # If BF16 type, +64 adds 128 bytes (correct, matches tStS+32 FP32 = 128 bytes)
        cute.printf("tStS ptr value: %d", tStS.iterator)
        cute.printf("tStS_P ptr (tStS+32): %d", tStS.iterator + tmem_p0_offset)
        cute.printf("tOrP ptr value: %d", tOrP.iterator)
        cute.printf("tOrP0 ptr (tOrP+64): %d", tOrP.iterator + p_offset)
        cute.printf("p_offset: %d", p_offset)

        pv_acc_shape = pv_thr.partition_shape_C(pv_mma_tiler[:2])
        tOtO = pv_thr.make_fragment_C(pv_acc_shape)
        o_cols = find_tmem_tensor_col_offset(tOtO)
        cute.printf("o_cols: %d", o_cols)

kernel = InspectKernel()
compiled = cute.compile(kernel, mQ, mK, mV, mC, stream)
compiled(mQ, mK, mV, mC, stream)
torch.cuda.synchronize()
