@triton.heuristics({'EVEN_K': lambda args: args['K'] % (args['BLOCK_K'] *
    args['SPLIT_K']) == 0})
@triton.autotune(configs=[triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256,
    'BLOCK_K': 32, 'SPLIT_K': 1}, num_stages=3, num_warps=8), triton.Config
    ({'BLOCK_M': 256, 'BLOCK_N': 128, 'BLOCK_K': 32, 'SPLIT_K': 1},
    num_stages=3, num_warps=8), triton.Config({'BLOCK_M': 256, 'BLOCK_N': 
    64, 'BLOCK_K': 32, 'SPLIT_K': 1}, num_stages=4, num_warps=4), triton.
    Config({'BLOCK_M': 64, 'BLOCK_N': 256, 'BLOCK_K': 32, 'SPLIT_K': 1},
    num_stages=4, num_warps=4), triton.Config({'BLOCK_M': 128, 'BLOCK_N': 
    128, 'BLOCK_K': 32, 'SPLIT_K': 1}, num_stages=4, num_warps=4), triton.
    Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 32, 'SPLIT_K': 1},
    num_stages=4, num_warps=4), triton.Config({'BLOCK_M': 64, 'BLOCK_N': 
    128, 'BLOCK_K': 32, 'SPLIT_K': 1}, num_stages=4, num_warps=4), triton.
    Config({'BLOCK_M': 128, 'BLOCK_N': 32, 'BLOCK_K': 32, 'SPLIT_K': 1},
    num_stages=4, num_warps=4), triton.Config({'BLOCK_M': 64, 'BLOCK_N': 32,
    'BLOCK_K': 32, 'SPLIT_K': 1}, num_stages=5, num_warps=2), triton.Config
    ({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 64, 'SPLIT_K': 1},
    num_stages=3, num_warps=8), triton.Config({'BLOCK_M': 256, 'BLOCK_N': 
    128, 'BLOCK_K': 64, 'SPLIT_K': 1}, num_stages=3, num_warps=8), triton.
    Config({'BLOCK_M': 256, 'BLOCK_N': 64, 'BLOCK_K': 64, 'SPLIT_K': 1},
    num_stages=4, num_warps=4), triton.Config({'BLOCK_M': 64, 'BLOCK_N': 
    256, 'BLOCK_K': 64, 'SPLIT_K': 1}, num_stages=2, num_warps=4), triton.
    Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 64, 'SPLIT_K': 1},
    num_stages=2, num_warps=4), triton.Config({'BLOCK_M': 128, 'BLOCK_N': 
    256, 'BLOCK_K': 64, 'SPLIT_K': 1}, num_stages=1, num_warps=8), triton.
    Config({'BLOCK_M': 256, 'BLOCK_N': 128, 'BLOCK_K': 64, 'SPLIT_K': 1},
    num_stages=1, num_warps=8), triton.Config({'BLOCK_M': 256, 'BLOCK_N': 
    64, 'BLOCK_K': 64, 'SPLIT_K': 1}, num_stages=1, num_warps=4), triton.
    Config({'BLOCK_M': 64, 'BLOCK_N': 256, 'BLOCK_K': 64, 'SPLIT_K': 1},
    num_stages=1, num_warps=4), triton.Config({'BLOCK_M': 128, 'BLOCK_N': 
    128, 'BLOCK_K': 64, 'SPLIT_K': 1}, num_stages=1, num_warps=4), triton.
    Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 64, 'SPLIT_K': 1},
    num_stages=4, num_warps=4), triton.Config({'BLOCK_M': 64, 'BLOCK_N': 
    128, 'BLOCK_K': 64, 'SPLIT_K': 1}, num_stages=4, num_warps=4), triton.
    Config({'BLOCK_M': 128, 'BLOCK_N': 32, 'BLOCK_K': 64, 'SPLIT_K': 1},
    num_stages=4, num_warps=4), triton.Config({'BLOCK_M': 64, 'BLOCK_N': 32,
    'BLOCK_K': 64, 'SPLIT_K': 1}, num_stages=5, num_warps=2), triton.Config
    ({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 64, 'SPLIT_K': 1},
    num_stages=1, num_warps=4), triton.Config({'BLOCK_M': 64, 'BLOCK_N': 
    128, 'BLOCK_K': 64, 'SPLIT_K': 1}, num_stages=1, num_warps=4), triton.
    Config({'BLOCK_M': 128, 'BLOCK_N': 32, 'BLOCK_K': 64, 'SPLIT_K': 1},
    num_stages=1, num_warps=4), triton.Config({'BLOCK_M': 64, 'BLOCK_N': 32,
    'BLOCK_K': 64, 'SPLIT_K': 1}, num_stages=1, num_warps=2)], key=['M',
    'N', 'K'])
@triton.jit
def _kernel(A, B, C, M, N, K, stride_am, stride_ak, stride_bk, stride_bn,
    stride_cm, stride_cn, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr, GROUP_M: tl.constexpr, SPLIT_K: tl.constexpr,
    EVEN_K: tl.constexpr, ACC_TYPE: tl.constexpr):
    pid = tl.program_id(0)
    pid_z = tl.program_id(1)
    bid = tl.program_id(2)
    grid_m = (M + BLOCK_M - 1) // BLOCK_M
    grid_n = (N + BLOCK_N - 1) // BLOCK_N
    width = GROUP_M * grid_n
    group_id = pid // width
    group_size = min(grid_m - group_id * GROUP_M, GROUP_M)
    pid_m = group_id * GROUP_M + pid % group_size
    pid_n = pid % width // group_size
    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M)
    rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N)
    rk = pid_z * BLOCK_K + tl.arange(0, BLOCK_K)
    A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak)
    B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn)
    A += bid * M * K
    B += bid * K * N
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE)
    for k in range(K, 0, -BLOCK_K * SPLIT_K):
        if EVEN_K:
            a = tl.load(A)
            b = tl.load(B)
        else:
            a = tl.load(A, mask=rk[None, :] < k, other=0.0)
            b = tl.load(B, mask=rk[:, None] < k, other=0.0)
        acc += tl.dot(a, b)
        A += BLOCK_K * SPLIT_K * stride_ak
        B += BLOCK_K * SPLIT_K * stride_bk
    acc = acc.to(C.dtype.element_ty)
    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    C = C + (rm[:, None] * stride_cm + rn[None, :] * stride_cn)
    C += bid * M * N
    mask = (rm < M)[:, None] & (rn < N)[None, :]
    if SPLIT_K == 1:
        tl.store(C, acc, mask=mask)
    else:
        tl.atomic_add(C, acc, mask=mask)
