Can someone please help me debug this code

[conquerer-wo]

so I am new to triton and have been trying to write a kernel that Fuses 2 matrix multiplications into one kernel, yet I keep receiving a PassManager::run failed error, could you please help me figure out where the problem is(btw im on triton version 3.3),
this is the code:

import torch
import triton
import triton.language as tl
@triton.jit
def matmul_kernel(

    a_ptr, b_ptr, c_ptr,
    a1_ptr,b1_ptr,c1_ptr,
    
    M, N, K,
    B, S,O,
    
    stride_a, stride_a1,  
    stride_b, stride_b1,  
    stride_ce, stride_c1,

    stride_am, stride_ak,  
    stride_bk, stride_bn,  
    stride_cm, stride_cn,
  

    BLOCK_SIZE_B: tl.constexpr, BLOCK_SIZE_S: tl.constexpr, BLOCK_SIZE_O: tl.constexpr,  #
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,  #
    GROUP_SIZE_M: tl.constexpr,
    
    GROUP_SIZE_B: tl.constexpr,  

):

pid = tl.program_id(axis=0)
num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
num_pid_in_group = GROUP_SIZE_M * num_pid_n
group_id = pid // num_pid_in_group
first_pid_m = group_id * GROUP_SIZE_M
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m

# -----------------------------------------------------------

tl.assume(pid_m >= 0)
tl.assume(pid_n >= 0)
tl.assume(stride_am > 0)
tl.assume(stride_ak > 0)
tl.assume(stride_bn > 0)
tl.assume(stride_bk > 0)
tl.assume(stride_cm > 0)
tl.assume(stride_cn > 0)


offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
offs_k = tl.arange(0, BLOCK_SIZE_K)
a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)

accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
    # Load the next block of A and B, generate a mask by checking the K dimension.
    # If it is out of bounds, set it to 0.
    a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
    b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
    #  accumulate along the K dimension.
    accumulator = tl.dot(a, b, accumulator)
    # Advance the ptrs to the next K block.
    a_ptrs += BLOCK_SIZE_K * stride_ak
    b_ptrs += BLOCK_SIZE_K * stride_bk


c = accumulator.to(tl.float16)


offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
tl.store(c_ptrs, c, mask=c_mask)

pid2 = tl.program_id(axis=1)
num_pid_b = tl.cdiv(B, BLOCK_SIZE_B)
num_pid_s = tl.cdiv(S, BLOCK_SIZE_S)
num_pid_in_groups = GROUP_SIZE_B * num_pid_s
group_ids = pid2 // num_pid_in_groups
first_pid_b = group_ids * GROUP_SIZE_B
group_size_b = min(num_pid_b - first_pid_b, GROUP_SIZE_B)
pid_b = first_pid_b + ((pid2 % num_pid_in_groups) % group_size_b)
pid_s = (pid2 % num_pid_in_groups) // group_size_b


offs_b = (pid_b * BLOCK_SIZE_B + tl.arange(0, BLOCK_SIZE_B)) % B
offs_s = (pid_s * BLOCK_SIZE_S + tl.arange(0, BLOCK_SIZE_S)) % S
offs_o = tl.arange(0, BLOCK_SIZE_O)
a1_ptrs = a1_ptr + (offs_b[:, None] * stride_a + offs_o[None, :] * stride_a1)
b1_ptrs = b1_ptr + (offs_o[:, None] * stride_b + offs_s[None, :] * stride_b)


acc = tl.zeros((BLOCK_SIZE_B, BLOCK_SIZE_S), dtype=tl.float32)
for o in range(0, tl.cdiv(O, BLOCK_SIZE_O)):
   
    i = tl.load(a1_ptrs, mask=offs_o[None, :] < O - o * BLOCK_SIZE_O, other=0.0)
    i2 = tl.load(b1_ptrs, mask=offs_o[:, None] < O - o * BLOCK_SIZE_O, other=0.0)
  
    acc = tl.dot(i, i2, acc)
   
    a1_ptrs += BLOCK_SIZE_O * stride_a1
    b1_ptrs += BLOCK_SIZE_O * stride_b1


c1 = acc.to(tl.float16)


offs_c = pid_b * BLOCK_SIZE_B + tl.arange(0, BLOCK_SIZE_B)
offs_c1 = pid_s * BLOCK_SIZE_S + tl.arange(0, BLOCK_SIZE_S)
c1_ptrs = c1_ptr + stride_ce * offs_c[:, None] + stride_c1 * offs_c1[None, :]
c1_mask = (offs_c[:, None] < B) & (offs_c1[None, :] < S)
tl.store(c1_ptrs, c1, mask=c1_mask)

def matmul(a: torch.Tensor, b: torch.Tensor, a2: torch.Tensor, b2: torch.Tensor):
M,K= a.shape
_, N = b.shape
B, O = a2.shape
_, S = b2.shape

a = a.cuda()
b = b.cuda()
a2 = a2.cuda()
b2 = b2.cuda()
c = torch.empty((M, N), device='cuda', dtype=torch.float16)
c1 = torch.empty((B, S), device='cuda', dtype=torch.float16)

meta = {
    'BLOCK_SIZE_B': 64, 'BLOCK_SIZE_S': 64, 'BLOCK_SIZE_O': 32,
    'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32,
    'GROUP_SIZE_M': 8,
    'GROUP_SIZE_B': 8
}
grid_size = (
    triton.cdiv(M, meta['BLOCK_SIZE_M']) * triton.cdiv(N, meta['BLOCK_SIZE_N']),
    triton.cdiv(B, meta['BLOCK_SIZE_B']) * triton.cdiv(S, meta['BLOCK_SIZE_S'])
)

matmul_kernel[grid_size](
    a, b, c,
    a2, b2, c1,
    M, N, K,
    B, S, O,
    a2.stride(0), a2.stride(1),
    b2.stride(0), b2.stride(1),
    c1.stride(0), c1.stride(1),
    a.stride(0), a.stride(1),
    b.stride(0), b.stride(1),
    c.stride(0), c.stride(1),

    **meta
)
return c,c1

f_t = torch.rand((4096,65536), dtype=torch.float16, device='cuda')
s_t = torch.rand((4096, 1024), dtype=torch.float16, device='cuda')
e_t = torch.rand((4096,65536), dtype=torch.float16, device='cuda')
b_t = torch.rand((65536,512), dtype=torch.float16, device='cuda')

c,c1 = matmul(f_t.t().contiguous(),s_t,e_t,b_t)
print(c.shape)
print(c1.shape)

[Jokeren]

There doesn't seem to be any problems on triton/main

[conquerer-wo]

so then do you think maybe its an error with the triton version or maybe because im working in google colab?