import torch
import triton
import triton.language as tl

# Define block sizes
BLOCK_SIZE_NB = 32
BLOCK_SIZE_MB = 32
BLOCK_SIZE_DB = 32

@triton.jit
def linear_log_sum_exp_kernel(
    E, C, LSE,
    E_stride_d, E_stride_n,
    C_stride_d, C_stride_v,
    N, D, V,
    BLOCK_SIZE_NB: tl.constexpr, BLOCK_SIZE_MB: tl.constexpr, BLOCK_SIZE_DB: tl.constexpr
):
    pid_n = tl.program_id(axis=0)  # Block ID along the N dimension
    pid_m = tl.program_id(axis=1)  # Block ID along the V dimension

    # Offsets for the block
    offsets_n = pid_n * BLOCK_SIZE_NB + tl.arange(0, BLOCK_SIZE_NB)
    offsets_m = pid_m * BLOCK_SIZE_MB + tl.arange(0, BLOCK_SIZE_MB)
    offsets_d = tl.arange(0, BLOCK_SIZE_DB)

    # Masks for bounds checking
    mask_n = offsets_n < N
    mask_m = offsets_m < V

    # Initialize accumulator for the block result
    acc = tl.zeros((BLOCK_SIZE_MB, BLOCK_SIZE_NB), dtype=tl.float32)

    # Loop over blocks in the D dimension
    for d_start in range(0, D, BLOCK_SIZE_DB):
        # Load a block of E (DB x NB)
        e_block = tl.load(
            E + d_start * E_stride_d + offsets_d[:, None] * E_stride_d + offsets_n[None, :] * E_stride_n,
            mask=(offsets_d[:, None] + d_start < D) & mask_n[None, :],
            other=0.0
        )

        # Load a block of C (DB x MB)
        c_block = tl.load(
            C + d_start * C_stride_d + offsets_d[:, None] * C_stride_d + offsets_m[None, :] * C_stride_v,
            mask=(offsets_d[:, None] + d_start < D) & mask_m[None, :],
            other=0.0
        )

        # Perform blockwise matrix multiplication
        acc += tl.dot(c_block, e_block)  

    # Compute log-sum-exp for this block
    max_val = tl.max(acc, axis=0)
    exp_acc = tl.exp(acc - max_val[None, :])
    lse_block = tl.log(tl.sum(exp_acc, axis=0)) + max_val

    # Update LSE with atomic operations (simulate atomic_max using atomic_add and exp)
    n_range = tl.arange(0, BLOCK_SIZE_NB)
    idx = offsets_n[n_range]
    old_val = tl.load(LSE + idx * LSE.stride(0), mask=idx < N, other=-float('inf'))
    new_val = lse_block[n_range]
    update = tl.where(new_val > old_val, tl.exp(new_val - old_val), 0.0)
    tl.atomic_add(LSE + idx * LSE.stride(0), update, priority=1, mask=idx < N)

def linear_log_sum_exp(E, C):
    D, N = E.shape
    _, V = C.shape

    # Initialize LSE
    LSE = torch.full((N,), float('-inf'), device=E.device, dtype=E.dtype)

    # Define grid size
    grid_n = (N + BLOCK_SIZE_NB - 1) // BLOCK_SIZE_NB
    grid_m = (V + BLOCK_SIZE_MB - 1) // BLOCK_SIZE_MB

    # Launch the Triton kernel
    linear_log_sum_exp_kernel[(grid_n, grid_m)](
        E, C, LSE,
        E.stride(0), E.stride(1),
        C.stride(0), C.stride(1),
        N, D, V,
        BLOCK_SIZE_NB=BLOCK_SIZE_NB,
        BLOCK_SIZE_MB=BLOCK_SIZE_MB,
        BLOCK_SIZE_DB=BLOCK_SIZE_DB
    )

    return LSE

# Example usage and benchmarking
import time

# Example matrices
D, N, V = 512, 1024, 50000
E = torch.randn((D, N), device='cuda', dtype=torch.float32)
C = torch.randn((D, V), device='cuda', dtype=torch.float32)

# Triton implementation
start_time = time.time()
lse_triton = linear_log_sum_exp(E, C)
triton_time = time.time() - start_time

# PyTorch implementation for comparison
start_time = time.time()
logits = torch.matmul(C.T, E)
lse_torch = torch.logsumexp(logits, dim=0)
torch_time = time.time() - start_time

# Compare outputs and speed
print("Triton Output (first 3 values):", lse_triton[:3].cpu().numpy())
print("Torch Output (first 3 values):", lse_torch[:3].cpu().numpy())
print("Are outputs close?", torch.allclose(lse_triton, lse_torch, atol=1e-5))
print(f"Triton Execution Time: {triton_time:.6f}s")
print(f"Torch Execution Time: {torch_time:.6f}s")


---------------------------------------------------------------------------
CompilationError                          Traceback (most recent call last)
Cell In[18], line 99
     97 # Triton implementation
     98 start_time = time.time()
---> 99 lse_triton = linear_log_sum_exp(E, C)
    100 triton_time = time.time() - start_time
    102 # PyTorch implementation for comparison

Cell In[18], line 77, in linear_log_sum_exp(E, C)
     74 grid_m = (V + BLOCK_SIZE_MB - 1) // BLOCK_SIZE_MB
     76 # Launch the Triton kernel
---> 77 linear_log_sum_exp_kernel[(grid_n, grid_m)](
     78     E, C, LSE,
     79     E.stride(0), E.stride(1),
     80     C.stride(0), C.stride(1),
     81     N, D, V,
     82     BLOCK_SIZE_NB=BLOCK_SIZE_NB,
     83     BLOCK_SIZE_MB=BLOCK_SIZE_MB,
     84     BLOCK_SIZE_DB=BLOCK_SIZE_DB
     85 )
     87 return LSE

File /opt/conda/lib/python3.10/site-packages/triton/runtime/jit.py:345, in KernelInterface.__getitem__.<locals>.<lambda>(*args, **kwargs)
    339 def __getitem__(self, grid) -> T:
    340     """
    341     A JIT function is launched with: fn[grid](*args, **kwargs).
    342     Hence JITFunction.__getitem__ returns a callable proxy that
    343     memorizes the grid.
    344     """
--> 345     return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/triton/runtime/jit.py:662, in JITFunction.run(self, grid, warmup, *args, **kwargs)
    660     # compile the kernel
    661     src = self.ASTSource(self, signature, constants, configs[0])
--> 662     kernel = self.compile(
    663         src,
    664         target=target,
    665         options=options.__dict__,
    666     )
    667     self.cache[device][key] = kernel
    669 # Check that used global values have not changed.

File /opt/conda/lib/python3.10/site-packages/triton/compiler/compiler.py:276, in compile(src, target, options)
    274 codegen_fns = backend.get_codegen_implementation()
    275 try:
--> 276     module = src.make_ir(options, codegen_fns, context)
    277 except Exception as e:
    278     filter_traceback(e)

File /opt/conda/lib/python3.10/site-packages/triton/compiler/compiler.py:113, in ASTSource.make_ir(self, options, codegen_fns, context)
    112 def make_ir(self, options, codegen_fns, context):
--> 113     return ast_to_ttir(self.fn, self, context=context, options=options, codegen_fns=codegen_fns)

CompilationError: at 49:10:

        # Perform blockwise matrix multiplication
        acc += tl.dot(c_block, e_block)  

    # Compute log-sum-exp for this block
    max_val = tl.max(acc, axis=0)
    exp_acc = tl.exp(acc - max_val[None, :])
    lse_block = tl.log(tl.sum(exp_acc, axis=0)) + max_val

    # Update LSE with atomic operations (simulate atomic_max using atomic_add and exp)
    n_range = tl.arange(0, BLOCK_SIZE_NB)
    idx = offsets_n[n_range]
          ^
ValueError('Did you forget to add @triton.jit ? (`_builder` argument must be provided outside of JIT functions.)')

please help me