from tilelang import tvm as tvm
import tilelang.testing
import tilelang.language as T

tilelang.testing.set_random_seed(0)


def chunk_scan_fwd(batch, seqlen, chunk_size, ngroups, nheads, headdim, dstate, block_M, block_N,
                   block_K, block_Dstate, num_stages, threads):
    dtype = "float16"
    accum_dtype = "float"
    nchunks = T.ceildiv(seqlen, chunk_size)
    p = 1.44269504

    @T.prim_func
    def main(cb: T.Buffer((batch, nchunks, ngroups, chunk_size, chunk_size), dtype), x: T.Buffer(
        (batch, seqlen, nheads, headdim), dtype), dt: T.Buffer(
            (batch, nheads, nchunks, chunk_size), dtype), dA_cumsum: T.Buffer(
                (batch, nheads, nchunks, chunk_size), dtype),
             C: T.Buffer((batch, seqlen, ngroups, dstate), dtype), prev_states: T.Buffer(
                 (batch, nchunks, nheads, headdim, dstate), dtype), D: T.Buffer(
                     (nheads), dtype), Output: T.Buffer((batch, seqlen, nheads, headdim), dtype)):
        with T.Kernel(
                nheads,
                T.ceildiv(chunk_size, block_M) * T.ceildiv(headdim, block_N),
                batch * nchunks,
                threads=threads) as (bz, bx, by):
            acc_o = T.alloc_fragment((block_M, block_N), accum_dtype)
            cb_shared = T.alloc_shared((block_M, block_K), dtype, scope="shared.dyn")
            cb_local = T.alloc_fragment((block_M, block_K), dtype)
            dA_cs_k_shared = T.alloc_shared((block_K), dtype, scope="shared")
            dA_cs_k_local = T.alloc_fragment((block_K), accum_dtype)
            dA_cs_m_local = T.alloc_fragment((block_M), accum_dtype)
            dt_shared = T.alloc_shared((block_K), dtype, scope="shared")
            dt_local = T.alloc_fragment((block_K), accum_dtype)
            x_shared = T.alloc_shared((block_K, block_N), dtype, scope="shared.dyn")
            dA_cs_m_shared = T.alloc_shared((block_M), dtype, scope="shared")
            scale_m_local = T.alloc_fragment((block_M), accum_dtype)
            C_shared = T.alloc_shared((block_M, block_Dstate), dtype)
            prev_state_shared = T.alloc_shared((block_N, block_Dstate), dtype)
            D_local = T.alloc_fragment((1), accum_dtype)
            x_residual_shared = T.alloc_shared((block_M, block_N), dtype, scope="shared.dyn")
            x_residual_local = T.alloc_fragment((block_M, block_N), accum_dtype)

            batch_idx = by % batch
            chunk_idx = by // batch
            # m: chunk_size
            # n : headdim
            m_idx = bx // T.ceildiv(headdim, block_N)
            n_idx = bx % T.ceildiv(headdim, block_N)

            T.copy(dA_cumsum[batch_idx, bz, chunk_idx, m_idx * block_M:(m_idx + 1) * block_M],
                   dA_cs_m_shared)
            T.copy(dA_cs_m_shared, dA_cs_m_local)
            T.clear(acc_o)

            for i in T.Parallel(block_M):
                scale_m_local[i] = T.exp2(dA_cs_m_local[i] * p)
            T.copy(
                C[batch_idx, chunk_idx * chunk_size + m_idx * block_M:chunk_idx * chunk_size +
                  (m_idx + 1) * block_M, bz // (nheads // ngroups), 0:block_Dstate], C_shared)
            T.copy(
                prev_states[batch_idx, chunk_idx, bz, n_idx * block_N:(n_idx + 1) * block_N,
                            0:block_Dstate], prev_state_shared)
            T.gemm(C_shared, prev_state_shared, acc_o, transpose_B=True)
            for i, j in T.Parallel(block_M, block_N):
                acc_o[i, j] *= scale_m_local[i]

            loop_range = T.ceildiv((m_idx + 1) * block_M, block_K)

            for k in T.Pipelined(loop_range, num_stages=num_stages):
                T.copy(
                    cb[batch_idx, chunk_idx, bz // (nheads // ngroups),
                       m_idx * block_M:(m_idx + 1) * block_M, k * block_K:(k + 1) * block_K],
                    cb_shared)
                T.copy(cb_shared, cb_local)
                T.copy(dA_cumsum[batch_idx, bz, chunk_idx, k * block_K:(k + 1) * block_K],
                       dA_cs_k_shared)
                T.copy(dA_cs_k_shared, dA_cs_k_local)
                for i, j in T.Parallel(block_M, block_K):
                    cb_local[i,
                             j] = cb_local[i,
                                           j] * T.exp2(dA_cs_m_local[i] * p - dA_cs_k_local[j] * p)
                T.copy(dt[batch_idx, bz, chunk_idx, k * block_K:(k + 1) * block_K], dt_shared)
                T.copy(dt_shared, dt_local)
                for i, j in T.Parallel(block_M, block_K):
                    cb_local[i, j] *= dt_local[j]
                for i, j in T.Parallel(block_M, block_K):
                    cb_local[i, j] = T.if_then_else(m_idx * block_M + i >= k * block_K + j,
                                                    cb_local[i, j], 0)
                T.copy(
                    x[batch_idx, chunk_idx * chunk_size + k * block_K:chunk_idx * chunk_size +
                      (k + 1) * block_K, bz, n_idx * block_N:(n_idx + 1) * block_N], x_shared)
                T.gemm(cb_local, x_shared, acc_o)

            D_local[0] = D[bz]
            T.copy(
                x[batch_idx, chunk_idx * chunk_size + m_idx * block_M:chunk_idx * chunk_size +
                  (m_idx + 1) * block_M, bz, n_idx * block_N:(n_idx + 1) * block_N],
                x_residual_shared)
            T.copy(x_residual_shared, x_residual_local)
            for i, j in T.Parallel(block_M, block_N):
                acc_o[i, j] += x_residual_local[i, j] * D_local[0]

            T.copy(
                acc_o,
                Output[batch_idx, chunk_idx * chunk_size + m_idx * block_M:chunk_idx * chunk_size +
                       (m_idx + 1) * block_M, bz, n_idx * block_N:(n_idx + 1) * block_N])

    return main


def run_chunk_scan(batch,
                   seqlen,
                   chunk_size,
                   ngroups,
                   nheads,
                   headdim,
                   dstate,
                   block_M,
                   block_N,
                   block_K,
                   block_Dstate,
                   num_stages=2,
                   threads=128):
    program = chunk_scan_fwd(batch, seqlen, chunk_size, ngroups, nheads, headdim, dstate, block_M,
                             block_N, block_K, block_Dstate, num_stages, threads)

    kernel = tilelang.compile(program, out_idx=[7])
    profiler = kernel.get_profiler()

    def ref_program(cb, x, dt, dA_cumsum, C, prev_states, D):
        import torch
        from einops import rearrange, repeat
        """
        Argument:
            cb: (batch, nchunks, ngroups, chunk_size, chunk_size)
            x: (batch, seqlen, nheads, headdim)
            dt: (batch, nheads, nchunks, chunk_size)
            dA_cumsum: (batch, nheads, nchunks, chunk_size)
            C: (batch, seqlen, ngroups, dstate)
            prev_states: (batch, nchunks, nheads, headdim, dstate)
            D: (nheads, headdim) or (nheads,)
            z: (batch, seqlen, nheads, headdim)
        Return:
            out: (batch, seqlen, nheads, headdim)
        """
        _, _, ngroups, _, _ = cb.shape
        batch, seqlen, nheads, headdim = x.shape
        # _, _, ngroups, dstate = B.shape
        # assert B.shape == (batch, seqlen, ngroups, dstate)
        _, _, nchunks, chunk_size = dt.shape
        assert seqlen == nchunks * chunk_size
        # assert C.shape == B.shape
        # B = repeat(B, "b l g d -> b l (g h) d", h=nheads // ngroups)
        C = repeat(C, "b l g d -> b l (g h) d", h=nheads // ngroups)
        cb = repeat(cb, "b c g l s -> b c (g h) l s", h=nheads // ngroups)
        # CB = torch.einsum("bclhn,bcshn->bchls", rearrange(C, "b (c l) h n -> b c l h n", c=nchunks),
        #                   rearrange(B, "b (c s) h n -> b c s h n", c=nchunks))
        # (batch, nheads, nchunks, chunksize, chunksize)
        dt_segment_sum = dA_cumsum[:, :, :, :, None] - dA_cumsum[:, :, :, None, :]
        decay = torch.exp(dt_segment_sum)
        scores_decay = cb * rearrange(decay, "b h c l s -> b c h l s")
        causal_mask = torch.tril(
            torch.ones(chunk_size, chunk_size, device=x.device, dtype=bool), diagonal=0)
        scores_decay = scores_decay.masked_fill(~causal_mask, 0)
        out = torch.einsum('bchls,bhcs,bcshp->bclhp', scores_decay.to(x.dtype), dt.to(x.dtype),
                           rearrange(x, "b (c s) h p -> b c s h p", c=nchunks))
        state_decay_out = torch.exp(rearrange(dA_cumsum, "b h c l -> b c l h 1"))
        out_prev = torch.einsum('bclhn,bchpn->bclhp',
                                rearrange(C, "b (c l) h n -> b c l h n", c=nchunks),
                                prev_states.to(C.dtype)) * state_decay_out
        out = out + out_prev
        out = rearrange(out, "b c l h p -> b (c l) h p")
        if D is not None:
            if D.dim() == 1:
                D = rearrange(D, "h -> h 1")
            out = out + x * D
        return out

    profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2, max_mismatched_ratio=0.05)


def chunk_state_fwd(batch,
                    seqlen,
                    chunk_size,
                    ngroups,
                    nheads,
                    headdim,
                    dstate,
                    block_M,
                    block_N,
                    block_K,
                    num_stages=2,
                    threads=128):
    dtype = "float16"
    accum_dtype = "float"
    nchunks = T.ceildiv(seqlen, chunk_size)
    p = 1.44269504

    @T.prim_func
    def main(B: T.Buffer((batch, seqlen, ngroups, dstate), dtype), x: T.Buffer(
        (batch, seqlen, nheads, headdim), dtype), dt: T.Buffer(
            (batch, nheads, nchunks, chunk_size), dtype), dA_cumsum: T.Buffer(
                (batch, nheads, nchunks, chunk_size), dtype), Output: T.Buffer(
                    (batch, nchunks, nheads, headdim, dstate), dtype)):
        with T.Kernel(
                nheads,
                T.ceildiv(headdim, block_M) * T.ceildiv(dstate, block_N),
                batch * nchunks,
                threads=threads) as (bz, bx, by):
            x_shared = T.alloc_shared((block_K, block_M), dtype)
            x_local = T.alloc_fragment((block_K, block_M), dtype)
            xt_local = T.alloc_fragment((block_M, block_K), dtype)
            B_shared = T.alloc_shared((block_K, block_N), dtype)
            dt_shared = T.alloc_shared((block_K), dtype)
            dA_cumsum_shared = T.alloc_shared((block_K), dtype)
            acc_o = T.alloc_fragment((block_M, block_N), accum_dtype)
            scale = T.alloc_fragment((block_K), accum_dtype)
            dA_cs_last = T.alloc_fragment((1), accum_dtype)
            dA_cumsum_local = T.alloc_fragment((block_K), accum_dtype)
            dt_local = T.alloc_fragment((block_K), accum_dtype)

            loop_range = T.ceildiv(chunk_size, block_K)

            batch_idx = by % batch
            chunk_idx = by // batch
            m_idx = bx // T.ceildiv(dstate, block_N)
            n_idx = bx % T.ceildiv(dstate, block_N)

            dA_cs_last[0] = dA_cumsum[batch_idx, bz, chunk_idx, chunk_size - 1]
            T.clear(acc_o)
            for k in T.Pipelined(loop_range, num_stages=num_stages):
                T.copy(
                    x[batch_idx, chunk_idx * chunk_size + k * block_K:chunk_idx * chunk_size +
                      (k + 1) * block_K, bz, m_idx * block_M:(m_idx + 1) * block_M], x_shared)
                T.copy(dA_cumsum[batch_idx, bz, chunk_idx, k * block_K:(k + 1) * block_K],
                       dA_cumsum_shared)
                T.copy(dt[batch_idx, bz, chunk_idx, k * block_K:(k + 1) * block_K], dt_shared)
                T.copy(dA_cumsum_shared, dA_cumsum_local)
                T.copy(dt_shared, dt_local)
                for i in T.Parallel(block_K):
                    scale[i] = T.exp2(dA_cs_last[0] * p - dA_cumsum_local[i] * p) * dt_local[i]
                T.copy(x_shared, x_local)
                for i, j in T.Parallel(block_M, block_K):
                    xt_local[i, j] = x_local[j, i] * scale[j]
                T.copy(
                    B[batch_idx, chunk_idx * chunk_size + k * block_K:chunk_idx * chunk_size +
                      (k + 1) * block_K, bz // (nheads // ngroups),
                      n_idx * block_N:(n_idx + 1) * block_N], B_shared)
                T.gemm(xt_local, B_shared, acc_o)
            T.copy(
                acc_o, Output[batch_idx, chunk_idx, bz, m_idx * block_M:(m_idx + 1) * block_M,
                              n_idx * block_N:(n_idx + 1) * block_N])

    return main


def run_chunk_state(batch,
                    seqlen,
                    chunk_size,
                    ngroups,
                    nheads,
                    headdim,
                    dstate,
                    block_M,
                    block_N,
                    block_K,
                    num_stages=2,
                    threads=128):
    program = chunk_state_fwd(batch, seqlen, chunk_size, ngroups, nheads, headdim, dstate, block_M,
                              block_N, block_K, num_stages, threads)

    kernel = tilelang.compile(program, out_idx=[4])
    profiler = kernel.get_profiler()

    def ref_program(B, x, dt, dA_cumsum):
        """
        Argument:
            B: (batch, seqlen, ngroups, headdim)
            x: (batch, seqlen, nheads, headdim)
            dt: (batch, nheads, nchunks, chunk_size)
            dA_cumsum: (batch, nheads, nchunks, chunk_size)
        Return:
            states: (batch, nchunks, nheads, headdim, dstate)
        """
        # Check constraints.
        import torch
        import torch.nn.functional as F
        from einops import rearrange, repeat

        batch, seqlen, nheads, headdim = x.shape
        dstate = B.shape[-1]
        _, _, nchunks, chunk_size = dt.shape
        assert seqlen <= nchunks * chunk_size
        assert x.shape == (batch, seqlen, nheads, headdim)
        assert dt.shape == (batch, nheads, nchunks, chunk_size)
        ngroups = B.shape[2]
        assert nheads % ngroups == 0
        assert B.shape == (batch, seqlen, ngroups, dstate)
        B = repeat(B, "b l g d -> b l (g h) d", h=nheads // ngroups)
        assert dA_cumsum.shape == (batch, nheads, nchunks, chunk_size)
        if seqlen < nchunks * chunk_size:
            x = F.pad(x, (0, 0, 0, 0, 0, nchunks * chunk_size - seqlen))
            B = F.pad(B, (0, 0, 0, 0, 0, nchunks * chunk_size - seqlen))
        x = rearrange(x, "b (c l) h p -> b c l h p", l=chunk_size)
        B = rearrange(B, "b (c l) ... -> b c l ...", l=chunk_size)
        decay_states = torch.exp((dA_cumsum[:, :, :, -1:] - dA_cumsum))
        return torch.einsum("bclhn,bhcl,bhcl,bclhp->bchpn", B.to(x.dtype), decay_states.to(x.dtype),
                            dt.to(x.dtype), x)

    profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2, max_mismatched_ratio=0.05)


def test_chunk_scan():
    run_chunk_scan(
        batch=8,
        seqlen=2048,
        chunk_size=256,
        ngroups=1,
        nheads=8,
        headdim=64,
        dstate=128,
        block_M=64,
        block_N=64,
        block_K=64,
        block_Dstate=128,
        num_stages=2,
        threads=128)


def test_chunk_state():
    run_chunk_state(
        batch=8,
        seqlen=2048,
        chunk_size=256,
        ngroups=1,
        nheads=8,
        headdim=64,
        dstate=128,
        block_M=64,
        block_N=64,
        block_K=64,
        num_stages=2,
        threads=128)


if __name__ == "__main__":
    tilelang.testing.main()
