import triton
import triton.language as tl
import torch

@triton.heuristics(
    {
        "EVEN_K": lambda args: args["K"] % args["BLOCK_K"] == 0,
        "GRID_MN": lambda args: triton.cdiv(args["M"], args["BLOCK_M"])
        * triton.cdiv(args["N"], args["BLOCK_N"]),
    }
)
@triton.jit
def fp8_gemm_kernel(
    a_ptr,          # Pointer to A [M, K], column-major, float8_e4m3fnuz
    b_ptr,          # Pointer to B [N, K], column-major, float8_e4m3fnuz
    c_ptr,          # Pointer to C [M, N], row-major, bfloat16
    a_scale_ptr,    # Pointer to a_scale [M, K // 128], column-major, float32
    b_scale_ptr,    # Pointer to b_scale [N // 128, K // 128], column-major, float32
    M, N, K,        # Matrix dimensions
    stride_a_m, stride_a_k,      # Strides for A
    stride_b_n, stride_b_k,      # Strides for B
    stride_c_m, stride_c_n,      # Strides for C
    stride_as_m, stride_as_k,    # Strides for a_scale
    stride_bs_n, stride_bs_k,    # Strides for b_scale
    BLOCK_M: tl.constexpr,       # Tile size for M
    BLOCK_N: tl.constexpr,       # Tile size for N
    BLOCK_K: tl.constexpr,       # Tile size for K
    GROUP_SIZE: tl.constexpr,
    EVEN_K: tl.constexpr,
    GRID_MN: tl.constexpr,
):
    # tl.assume(stride_a_m > 0)
    # tl.assume(stride_a_k > 0)
    # tl.assume(stride_b_k > 0)
    # tl.assume(stride_b_n > 0)
    # tl.assume(stride_c_m > 0)
    # tl.assume(stride_c_n > 0)
    # tl.assume(stride_as_m > 0)
    # tl.assume(stride_as_k > 0)
    # tl.assume(stride_bs_k > 0)
    # tl.assume(stride_bs_n > 0)

    # GROUP_SIZE = 8
    # PID = tl.program_id(axis=0) 
    # # defining the size of groups
    # num_PID_along_M = tl.cdiv(M, BLOCK_M) # the number of blocks along M dimension
    # num_PID_along_N = tl.cdiv(N, BLOCK_N) # the number of blocks along N dimension
    # num_PID_in_group = GROUP_SIZE * num_PID_along_N
    # # figurinig out which group this PID is in
    # group_id = PID // num_PID_in_group
    # # tells us which row to start at for this group
    # first_PID_in_group_along_M = group_id * GROUP_SIZE 
    # # this is usually equal to GROUP_SIZE; the alternative case happens when we're at edge of the tensor and 
    # #  its dimensions don't cleanly divde into GROUP_SIZE # TODO is this true?
    # group_size_adj = min(num_PID_along_M - first_PID_in_group_along_M, GROUP_SIZE) 
    # # this is the bulk of the actual mapping of PIDs to group-major ordering
    # pid_m = first_PID_in_group_along_M + ((PID % num_PID_in_group) % group_size_adj)
    #     # (PID % num_PID_in_group) puts the current program id into the context of a group
    #     # (first_PID_in_group_along_m + ...) shifts the PID into the correct group
    #     # (... % group_size_adj) removes the column component to get us onto the correct row
    # pid_n = (PID % num_PID_in_group) // group_size_adj

    NUM_XCDS: tl.constexpr = 1

    # Map program ids `pid` to the block of C it should compute.
    # This is done in a grouped ordering to promote L2 data reuse.
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)

    ## pid remapping on xcds
    # Number of pids per XCD in the new arrangement
    pids_per_xcd = (GRID_MN + NUM_XCDS - 1) // NUM_XCDS
    # When GRID_MN cannot divide NUM_XCDS, some xcds will have
    # pids_per_xcd pids, the other will have pids_per_xcd - 1 pids.
    # We calculate the number of xcds that have pids_per_xcd pids as
    # tall_xcds
    tall_xcds = GRID_MN % NUM_XCDS
    tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds
    # Compute current XCD and local pid within the XCD
    xcd = pid % NUM_XCDS
    local_pid = pid // NUM_XCDS
    # Calculate new pid based on the new grouping
    # Note that we need to consider the following two cases:
    # 1. the current pid is on a tall xcd
    # 2. the current pid is on a short xcd
    if xcd < tall_xcds:
        pid = xcd * pids_per_xcd + local_pid
    else:
        pid = (
            tall_xcds * pids_per_xcd
            + (xcd - tall_xcds) * (pids_per_xcd - 1)
            + local_pid
        )

    if GROUP_SIZE == 1:
        pid_m = pid // num_pid_n
        pid_n = pid % num_pid_n
    else:
        num_pid_in_group = GROUP_SIZE * num_pid_n
        group_id = pid // num_pid_in_group
        first_pid_m = group_id * GROUP_SIZE
        group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE)
        pid_m = first_pid_m + (pid % group_size_m)
        pid_n = (pid % num_pid_in_group) // group_size_m


    # tl.assume(pid_m > 0)
    # tl.assume(pid_n > 0)
    
    # # Get block indices in the 2D grid
    # pid_m = tl.program_id(axis=0)  # Block index for M dimension
    # pid_n = tl.program_id(axis=1)  # Block index for N dimension

    # Compute offsets for the C tile
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) % M # Row indices for C
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) % N  # Column indices for C

    # Initialize accumulator for the C tile in FP32
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    it_k = tl.arange(0, BLOCK_K)
    scale_n_idx = (offs_n // 128)

    broadcasted_offs_m = offs_m[:, None]
    broadcasted_offs_n_r = offs_n[:, None]
    broadcasted_offs_n_c = offs_n[None, :]

    # K_scale = tl.cdiv(K, 128)
    # N_scale = tl.cdiv(N, 128)

    # Loop over K dimension in chunks of BLOCK_K
    # for k in tl.range(0, K, BLOCK_K, num_stages=3): # 1(none) , 2 or 3
    for k in range(0, K, BLOCK_K): # 2 or 3

        # Compute K offsets
        offs_k = k + it_k

        broadcasted_offs_k = offs_k[None, :]
        


        # Load tiles from A and B (column-major)
        # a_mask = (offs_m[:, None] < M) & (offs_k[None, :] < K)
        # b_mask = (offs_n[:, None] < N) & (offs_k[None, :] < K)


        if EVEN_K:
            a = tl.load(a_ptr + broadcasted_offs_m * stride_a_m + broadcasted_offs_k * stride_a_k)
            b = tl.load(b_ptr + broadcasted_offs_n_r * stride_b_n + broadcasted_offs_k * stride_b_k)
        else:
            # a_mask = (offs_m[:, None] < M) & (offs_k[None, :] < K)
            # b_mask = (offs_n[:, None] < N) & (offs_k[None, :] < K)
            a = tl.load(a_ptr + broadcasted_offs_m * stride_a_m + broadcasted_offs_k * stride_a_k, mask=(offs_m[:, None] < M) & (offs_k[None, :] < K), other=0.0)
            b = tl.load(b_ptr + broadcasted_offs_n_r * stride_b_n + broadcasted_offs_k * stride_b_k, mask=(offs_n[:, None] < N) & (offs_k[None, :] < K), other=0.0)
        
        # a = tl.load(
        #     a_ptr + offs_m[:, None] * stride_a_m + offs_k[None, :] * stride_a_k,
        #     mask=a_mask,
        #     other=0.0
        # )
        # b = tl.load(
        #     b_ptr + offs_n[:, None] * stride_b_n + offs_k[None, :] * stride_b_k,
        #     mask=b_mask,
        #     other=0.0
        # )

        # Compute scale indices for a_scale and b_scale
        scale_k_idx = k // 128  # Index into K // 128 dimension
        # a_scale_mask = (offs_m < M) & (scale_k_idx < K_scale)
        # b_scale_mask = ((offs_n // 128) < N_scale) & (scale_k_idx < K_scale)

        # Load a_scale [M, K // 128] for 1x128 chunks
        a_scale = tl.load(
            a_scale_ptr + offs_m * stride_as_m + scale_k_idx * stride_as_k,
            # mask=a_scale_mask,
            # other=1.0
        )
        # Load b_scale [N // 128, K // 128] for 128x128 chunks
        b_scale = tl.load(
            b_scale_ptr + scale_n_idx * stride_bs_n + scale_k_idx * stride_bs_k,
            # mask=b_scale_mask,
            # other=1.0
        )

        # Apply scaling block-wise
        # a_scaled = a * a_scale[:, None]
        # b_scaled = b * b_scale[:, None]
        # b_scale = b_scale[scale_k_idx]

        # Compute dot product
        # acc = tl.dot(a * a_scale[:, None], tl.trans(b * b_scale[:, None]), acc=acc)

        # Compute block result
        acc += tl.dot(a, tl.trans(b), out_dtype=tl.float32, input_precision="ieee") * a_scale[:, None] * b_scale[None, :]
    
    # Write results to C (row-major)
    # c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(
        c_ptr + broadcasted_offs_m * stride_c_m + broadcasted_offs_n_c * stride_c_n,
        acc.to(tl.bfloat16),
        mask=(broadcasted_offs_m < M) & (broadcasted_offs_n_c < N)
    )

def fp8_blockwise_gemm(
    a: torch.Tensor,          # [M, K], float8_e4m3fnuz, column-major
    b: torch.Tensor,          # [N, K], float8_e4m3fnuz, column-major
    a_scale: torch.Tensor,    # [M, K // 128], float32, column-major
    b_scale: torch.Tensor,    # [N // 128, K // 128], float32, column-major
    c: torch.Tensor           # [M, N], bfloat16, row-major
):
    assert a.shape[1] == b.shape[1], "Incompatible K dimension"
    assert a.is_cuda and b.is_cuda and a_scale.is_cuda and b_scale.is_cuda and c.is_cuda, "Tensors must be on GPU"
    assert a.dtype == torch.float8_e4m3fnuz, "A must be FP8"
    assert b.dtype == torch.float8_e4m3fnuz, "B must be FP8"
    assert a_scale.dtype == torch.float32, "a_scale must be FP32"
    assert b_scale.dtype == torch.float32, "b_scale must be FP32"
    assert c.dtype == torch.bfloat16, "C must be BF16"

    M, N = c.shape
    _, K = a.shape

    BLOCK_M = 128
    BLOCK_N = 128
    BLOCK_K = 128
    GROUP_SIZE = 4
    num_warps = 4
    num_stages = 2
    kpack = 2 # only for gemm kernel
    # waves_per_eu = 2
    # matrix_instr_nonkdim = 16

    b = b.contiguous()
    b_scale = b_scale.contiguous()

    print(a.is_contiguous(), a_scale.is_contiguous())
    print(b.is_contiguous(), b_scale.is_contiguous())

    # 2D grid
    # grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
    # 1D grid
    grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N), 1)
   
    # Launch kernel
    fp8_gemm_kernel[grid](
        a_ptr=a,
        b_ptr=b,
        c_ptr=c,
        a_scale_ptr=a_scale,
        b_scale_ptr=b_scale,
        M=M, N=N, K=K,
        stride_a_m=a.stride(0), stride_a_k=a.stride(1),
        stride_b_n=b.stride(0), stride_b_k=b.stride(1),
        stride_c_m=c.stride(0), stride_c_n=c.stride(1),
        stride_as_m=a_scale.stride(0), stride_as_k=a_scale.stride(1),
        stride_bs_n=b_scale.stride(0), stride_bs_k=b_scale.stride(1),
        BLOCK_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
        BLOCK_K=BLOCK_K,
        GROUP_SIZE=GROUP_SIZE,
        # waves_per_eu=waves_per_eu,
        kpack=kpack,
        # matrix_instr_nonkdim=matrix_instr_nonkdim,
        num_warps=num_warps,
        num_stages=num_stages,

    )
    return c


# SPDX-License-Identifier: MIT
# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.

# import torch
# import triton
# import triton.language as tl
# from typing import Optional


# # TODO Move this to a common folder. Will need to add future arch list
# def get_arch():
#     return triton.runtime.driver.active.get_current_target().arch


# @triton.heuristics(
#     {
#         "EVEN_K": lambda args: args["K"] % args["BLOCK_SIZE_K"] == 0,
#         "GRID_MN": lambda args: triton.cdiv(args["M"], args["BLOCK_SIZE_M"])
#         * triton.cdiv(args["N"], args["BLOCK_SIZE_N"]),
#     }
# )
# @triton.jit
# def _gemm_a8w8_blockscale_kernel(
#     # Pointers to matrices
#     a_ptr,
#     b_ptr,
#     c_ptr,
#     a_scale_ptr,
#     b_scale_ptr,
#     # Matrix dimensions
#     M,
#     N,
#     K,
#     # The stride variables represent how much to increase the ptr by when
#     # moving by 1 element in a particular dimension. E.g. `stride_am` is
#     # how much to increase `a_ptr` by to get the element one row down
#     # (A has M rows).
#     stride_am,
#     stride_ak,
#     stride_bk,
#     stride_bn,
#     stride_cm,
#     stride_cn,
#     stride_ascale_m,
#     stride_ascale_k,
#     stride_bscale_k,
#     stride_bscale_n,
#     # Meta-parameters
#     GROUP_K: tl.constexpr,
#     GROUP_N: tl.constexpr,
#     BLOCK_SIZE_M: tl.constexpr,
#     BLOCK_SIZE_N: tl.constexpr,
#     BLOCK_SIZE_K: tl.constexpr,
#     GROUP_SIZE_M: tl.constexpr,
#     EVEN_K: tl.constexpr,
#     GRID_MN: tl.constexpr,
# ):
#     """
#     Note: this is Triton jited function and not meant to be called directly. Call gemm_a8w8_blockscale function
#     below

#     Computes the 8 bit matmul C = A x B using the block-scale quantization approach.

#     Key parameters:
#     - A: Matrix A with shape (M, K).
#     - B: Matrix B with shape (K, N).
#     - C: Matrix C with shape (M, N).
#     - A_scale: Scale tensor for A with shape (M, *scale_k).
#     - B_scale: Scale tensor for B with shape (*scale_k, **scale_n).

#     *scale_k = (K + GROUP_K - 1) // GROUP_K
#     **scale_n = (N + GROUP_N - 1) // GROUP_N
#     """

#     NUM_XCDS: tl.constexpr = 8

#     # tl.assume(stride_am > 0)
#     # tl.assume(stride_ak > 0)
#     # tl.assume(stride_bk > 0)
#     # tl.assume(stride_bn > 0)
#     # tl.assume(stride_cm > 0)
#     # tl.assume(stride_cn > 0)
#     # tl.assume(stride_ascale_m > 0)
#     # tl.assume(stride_ascale_k > 0)
#     # tl.assume(stride_bscale_k > 0)
#     # tl.assume(stride_bscale_n > 0)

#     # -----------------------------------------------------------
#     # Map program ids `pid` to the block of C it should compute.
#     # This is done in a grouped ordering to promote L2 data reuse.
#     pid = tl.program_id(axis=0)
#     num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
#     num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)

#     ## pid remapping on xcds
#     # Number of pids per XCD in the new arrangement
#     pids_per_xcd = (GRID_MN + NUM_XCDS - 1) // NUM_XCDS
#     # When GRID_MN cannot divide NUM_XCDS, some xcds will have
#     # pids_per_xcd pids, the other will have pids_per_xcd - 1 pids.
#     # We calculate the number of xcds that have pids_per_xcd pids as
#     # tall_xcds
#     tall_xcds = GRID_MN % NUM_XCDS
#     tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds
#     # Compute current XCD and local pid within the XCD
#     xcd = pid % NUM_XCDS
#     local_pid = pid // NUM_XCDS
#     # Calculate new pid based on the new grouping
#     # Note that we need to consider the following two cases:
#     # 1. the current pid is on a tall xcd
#     # 2. the current pid is on a short xcd
#     if xcd < tall_xcds:
#         pid = xcd * pids_per_xcd + local_pid
#     else:
#         pid = (
#             tall_xcds * pids_per_xcd
#             + (xcd - tall_xcds) * (pids_per_xcd - 1)
#             + local_pid
#         )

#     if GROUP_SIZE_M == 1:
#         pid_m = pid // num_pid_n
#         pid_n = pid % num_pid_n
#     else:
#         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 % group_size_m)
#         pid_n = (pid % num_pid_in_group) // group_size_m

#     # tl.assume(pid_m > 0)
#     # tl.assume(pid_n > 0)

#     # Create pointers for first block of A and B input matrices
#     offs_k = tl.arange(0, BLOCK_SIZE_K)
#     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
#     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)

#     # Create pointers for the scales
#     k_start = 0
#     offs_ks = k_start // GROUP_K
#     a_scale_ptrs = a_scale_ptr + offs_am * stride_ascale_m + offs_ks * stride_ascale_k
#     offs_bsn = offs_bn // GROUP_N
#     b_scale_ptrs = b_scale_ptr + offs_ks * stride_bscale_k + offs_bsn * stride_bscale_n

#     acc_dtype = tl.float32 if c_ptr.type.element_ty != tl.int8 else tl.int32
#     accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=acc_dtype)

#     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.
#         if EVEN_K:
#             a = tl.load(a_ptrs)
#             b = tl.load(b_ptrs)
#         else:
#             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)

#         a_scale = tl.load(a_scale_ptrs)
#         b_scale = tl.load(b_scale_ptrs)

#         # Perform dot operation and apply scale
#         accumulator += (
#             tl.dot(a, b, input_precision="ieee") * a_scale[:, None] * b_scale[None, :]
#         )

#         # Advance the ptrs to the next K block.
#         a_ptrs += BLOCK_SIZE_K * stride_ak
#         b_ptrs += BLOCK_SIZE_K * stride_bk

#         k_cur = k * BLOCK_SIZE_K // GROUP_K
#         k_nxt = (k + 1) * BLOCK_SIZE_K // GROUP_K
#         offs_ks = k_nxt - k_cur
#         a_scale_ptrs += offs_ks * stride_ascale_k
#         b_scale_ptrs += offs_ks * stride_bscale_k

#     c = accumulator.to(c_ptr.type.element_ty)

#     # Write back the block of the output matrix C with masks.
#     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)


# def gemm_a8w8_blockscale(
#     x: torch.Tensor,
#     w: torch.Tensor,
#     x_scale: torch.Tensor,
#     w_scale: torch.Tensor,
#     dtype: Optional[float] = torch.bfloat16,
#     y: Optional[torch.Tensor] = None,
# ):
#     """
#     Computes the 8 bit matmul Y = X x WT using the block-scale quantization approach.

#     Key parameters:
#     - X: Matrix X with shape (M, K).
#     - W: Matrix W with shape (N, K).
#     - X_scale: Scale tensor for X with shape (M, *scale_k).
#     - W_scale: Scale tensor for W with shape (**scale_n, *scale_k).
#     - Y: Output Matrix Y with shape (M, K). If this is none, then it's created by this API and returned as output

#     Returns:
#     - Y: The output matrix with shape (M, N).

#     *scale_k = (K + scale_block_size_k - 1) // scale_block_size_k
#     **scale_n = (N + scale_block_size_n - 1) // scale_block_size_n
#     """

#     # Transpose w and w_scale
#     w = w.T
#     w_scale = w_scale.T

#     M, K = x.shape
#     K, N = w.shape

#     # Scale block sizes
#     # TODO: need a better way to pass scale block sizes around
#     GROUP_K = triton.next_power_of_2(triton.cdiv(K, w_scale.shape[0]))
#     GROUP_N = triton.next_power_of_2(triton.cdiv(N, w_scale.shape[1]))

#     # Check constraints.
#     assert x.shape[1] == w.shape[0], "Incompatible dimensions!!!"

#     # if y is None:
#     #     y = torch.empty((M, N), dtype=dtype, device=x.device)

#     BLOCK_SIZE_M = 128
#     BLOCK_SIZE_N = 128
#     BLOCK_SIZE_K = 128
#     GROUP_SIZE_M = 4
#     waves_per_eu = 2
#     kpack = 1 if get_arch() in ("gfx950") else 2
#     matrix_instr_nonkdim = 16
#     num_warps = 4
#     num_stages = 2

#     grid = (triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N),)

#     _gemm_a8w8_blockscale_kernel[grid](
#         x,
#         w,
#         y,
#         x_scale,
#         w_scale,
#         M,
#         N,
#         K,
#         x.stride(0),
#         x.stride(1),
#         w.stride(0),
#         w.stride(1),
#         y.stride(0),
#         y.stride(1),
#         x_scale.stride(0),
#         x_scale.stride(1),
#         w_scale.stride(0),
#         w_scale.stride(1),
#         GROUP_K,
#         GROUP_N,
#         BLOCK_SIZE_M,
#         BLOCK_SIZE_N,
#         BLOCK_SIZE_K,
#         GROUP_SIZE_M,
#         # waves_per_eu=waves_per_eu,
#         kpack=kpack,
#         # matrix_instr_nonkdim=matrix_instr_nonkdim,
#         num_warps=num_warps,
#         num_stages=num_stages,
#     )

#     return y


from task import input_t, output_t

def custom_kernel(data: input_t) -> output_t:
    """
    Efficient Triton implementation of FP8 GEMM for AMD MI300.
    
    Args:
        data: Tuple of (a, b, a_scale, b_scale, c)
    
    Returns:
        c: torch.Tensor[bfloat16] of shape [m, n]
    """
    a, b, a_scale, b_scale, c = data
    return fp8_blockwise_gemm(a, b, a_scale, b_scale, c)


# # autotune_configs = [
# #     triton.Config({'BLOCK_M': 64, 'BLOCK_N': 256, 'BLOCK_K': 32, 'GROUP_SIZE': 8}, num_stages=4, num_warps=4),
# #     triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 32, 'GROUP_SIZE': 8}, num_stages=4, num_warps=4),
# #     triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 32, 'GROUP_SIZE': 8}, num_stages=4, num_warps=4),
# #     triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128, 'BLOCK_K': 32, 'GROUP_SIZE': 8}, num_stages=4, num_warps=4),
# #     triton.Config({'BLOCK_M': 128, 'BLOCK_N': 32, 'BLOCK_K': 32, 'GROUP_SIZE': 8}, num_stages=4, num_warps=4),
# #     triton.Config({'BLOCK_M': 64, 'BLOCK_N': 32, 'BLOCK_K': 32, 'GROUP_SIZE': 8}, num_stages=5, num_warps=2),
# #     triton.Config({'BLOCK_M': 32, 'BLOCK_N': 64, 'BLOCK_K': 32, 'GROUP_SIZE': 8}, num_stages=5, num_warps=2)
# # ]

# # @triton.autotune(configs = autotune_configs, key=['M', 'N', 'K'])
python eval.py test /tmp/tmpp20yam4b
False True
True True
False False
True True
False False
True True
False False
True True
False False
True True
False False
True True
False False
True True
False False
True True
False False
True True
False False
True True
False False
True True
