Comparar texto

Encontre a diferença entre dois arquivos de texto

Editor live

Ocultar espaços em branco

Recolher inalteradas

Sem quebra de linha

Layout

Nível de detalhe

Estilos de texto

Realce de sintaxe

Diffchecker Desktop A maneira mais segura de usar o Diffchecker. Obtenha o aplicativo Diffchecker Desktop: seus diffs nunca saem do seu computador!Obter Desktop

Comparing TLX between persistent and CLC

Criado há 8 mesesO diff nunca expira

35 remoções

Linhas
Total
Removido

Caracteres
Total
Removido

Para continuar usando este recurso, atualize para Diffchecker Pro Ver preços

158 linhas

54 adições

Linhas
Total
Adicionado

Caracteres
Total
Adicionado

Para continuar usando este recurso, atualize para Diffchecker Pro Ver preços

191 linhas

@triton.jit

def matmul_kernel_tma_ws_blackwell_clc(a_desc, b_desc, c_desc, M, N, K, BLOCK_SIZE_M: tl.constexpr,

def matmul_kernel_tma_ws_blackwell_persistent(a_desc, b_desc, c_desc, M, N, K, BLOCK_SIZE_M: tl.constexpr,

BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, #

GROUP_SIZE_M: tl.constexpr, #

NUM_SMEM_BUFFERS: tl.constexpr, #

NUM_TMEM_BUFFERS: tl.constexpr, #

NUM_SMS: tl.constexpr, #

NUM_CLC_STAGES: tl.constexpr, #

EPILOGUE_SUBTILE: tl.constexpr, #

# allocate NUM_SMEM_BUFFERS buffers

buffers_A = tlx.local_alloc((BLOCK_SIZE_M, BLOCK_SIZE_K), tl.float16, NUM_SMEM_BUFFERS)

buffers_B = tlx.local_alloc((BLOCK_SIZE_K, BLOCK_SIZE_N), tl.float16, NUM_SMEM_BUFFERS)

# use multiple TMEM buffers to overlap MMA and epilogue

tmem_buffers = tlx.local_alloc((BLOCK_SIZE_M, BLOCK_SIZE_N), tl.float32, NUM_TMEM_BUFFERS, tlx.storage_kind.tmem)

# allocate barriers

smem_empty_bars = tlx.alloc_barriers(num_barriers=NUM_SMEM_BUFFERS, arrive_count=1)

smem_full_bars = tlx.alloc_barriers(num_barriers=NUM_SMEM_BUFFERS, arrive_count=1)

tmem_full_bars = tlx.alloc_barriers(num_barriers=NUM_TMEM_BUFFERS, arrive_count=1)

tmem_empty_bars = tlx.alloc_barriers(num_barriers=NUM_TMEM_BUFFERS, arrive_count=1)

clc_context = tlx.clc_create_context(NUM_CLC_STAGES, 3)

with tlx.async_tasks():

with tlx.async_task("default"): # epilogue consumer

# common code duplicated for each region to avoid SMEM overhead

start_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

num_tiles = num_pid_m * num_pid_n

k_tiles = tl.cdiv(K, BLOCK_SIZE_K)

# end of common code

tmem_read_phase = 0

cur_tmem_buf = 0

for tile_id in range(start_pid, num_tiles, NUM_SMS):

tile_id = start_pid

clc_phase_producer = 1

clc_phase_consumer = 0

clc_buf = 0

while tile_id != -1:

clc_buf = clc_buf % NUM_CLC_STAGES

# Debug prints

# if tlx.thread_id(axis=0) == 0:

# tl.device_print("Default WG Processing CtaID", tile_id)

# producer

tlx.clc_producer(clc_context, clc_buf, clc_phase_producer)

# clc_phase_producer ^= 1

clc_phase_producer = clc_phase_producer ^ (clc_buf == (NUM_CLC_STAGES - 1))

pid_m, pid_n = _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M)

offs_am = pid_m * BLOCK_SIZE_M

offs_bn = pid_n * BLOCK_SIZE_N

tlx.barrier_wait(tmem_full_bars[cur_tmem_buf], tmem_read_phase)

# flip phase at the end of a round of using TMEM barriers

tmem_read_phase = tmem_read_phase ^ (cur_tmem_buf == NUM_TMEM_BUFFERS - 1)

# load the result from TMEM to registers

acc_tmem = tmem_buffers[cur_tmem_buf]

if EPILOGUE_SUBTILE:

# We load/store the result half by half to reduce SMEM pressure

acc_tmem_subslice1 = tlx.subslice(acc_tmem, 0, BLOCK_SIZE_N // 2)

result = tlx.local_load(acc_tmem_subslice1)

c = result.to(tl.float16)

c_desc.store([offs_am, offs_bn], c)

acc_tmem_subslice2 = tlx.subslice(acc_tmem, BLOCK_SIZE_N // 2, BLOCK_SIZE_N // 2)

result = tlx.local_load(acc_tmem_subslice2)

c = result.to(tl.float16)

c_desc.store([offs_am, offs_bn + BLOCK_SIZE_N // 2], c)

else:

result = tlx.local_load(acc_tmem)

c = result.to(tl.float16)

c_desc.store([offs_am, offs_bn], c)

# done storing this buffer, signal MMA consumer to resume writing to it

tlx.barrier_arrive(tmem_empty_bars[cur_tmem_buf], 1)

cur_tmem_buf = (cur_tmem_buf + 1) % NUM_TMEM_BUFFERS

tile_id = tlx.clc_consumer(clc_context, clc_buf, clc_phase_consumer)

# clc_phase_consumer ^= 1

clc_phase_consumer = clc_phase_consumer ^ (clc_buf == (NUM_CLC_STAGES - 1))

clc_buf += 1

# Debug-only: verifying that CLC steals workloads successfully

# if tlx.thread_id(axis=0) == 0:

# tl.device_print("Extracted CtaID", tile_id)

with tlx.async_task(num_warps=1, num_regs=232): # MMA consumer

# common code duplicated for each region to avoid SMEM overhead

start_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

num_tiles = num_pid_m * num_pid_n

k_tiles = tl.cdiv(K, BLOCK_SIZE_K)

# end of common code

dot_phase = 0 # the current phase of dot op

tmem_write_phase = 1 # sync between epilogue consumer and MMA consumer

cur_tmem_buf = 0

processed_k_iters = 0

for tile_id in range(start_pid, num_tiles, NUM_SMS):

tile_id = start_pid

clc_phase = 0

clc_buf = 0

while tile_id != -1:

clc_buf = clc_buf % NUM_CLC_STAGES

pid_m, pid_n = _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M)

offs_am = pid_m * BLOCK_SIZE_M

offs_bn = pid_n * BLOCK_SIZE_N

# wait epilogue consumer to be done with the buffer before reusing it

tlx.barrier_wait(tmem_empty_bars[cur_tmem_buf], tmem_write_phase)

# flip phase at the end of a round of using TMEM barriers

tmem_write_phase = tmem_write_phase ^ (cur_tmem_buf == NUM_TMEM_BUFFERS - 1)

# now iterate along K to compute result for the block

for k in range(0, k_tiles):

# processed_k_iters + k means we use the immediate next buffer slot of tile_id x when we start tile_id x+1

buf = (processed_k_iters + k) % NUM_SMEM_BUFFERS

# wait for current phase(round) of load for this buf

tlx.barrier_wait(smem_full_bars[buf], dot_phase)

# buffer is now ready with loaded data, tlx.async_dot will signal `mBarrier` when done

tlx.async_dot(

tlx.async_dot(buffers_A[buf], buffers_B[buf], tmem_buffers[cur_tmem_buf], use_acc=k > 0,

buffers_A[buf],

mBarriers=[smem_empty_bars[buf]], out_dtype=tl.float32)

buffers_B[buf],

tmem_buffers[cur_tmem_buf],

use_acc=k > 0,

mBarriers=[smem_empty_bars[buf]],

out_dtype=tl.float32,

)

# flip phase at the end of a round

dot_phase = dot_phase ^ (buf == NUM_SMEM_BUFFERS - 1)

# wait for last mma to complete

last_buf = (processed_k_iters + k_tiles - 1) % NUM_SMEM_BUFFERS

# in case phase was flipped, we should use the phase value when dot op was issued

last_dot_phase = dot_phase ^ (last_buf == NUM_SMEM_BUFFERS - 1)

tlx.barrier_wait(smem_empty_bars[last_buf], last_dot_phase)

# done filling this buffer, signal epilogue consumer

tlx.barrier_arrive(tmem_full_bars[cur_tmem_buf], 1)

# possibly enter next iteration (next tile) without waiting for epilogue

cur_tmem_buf = (cur_tmem_buf + 1) % NUM_TMEM_BUFFERS

processed_k_iters += k_tiles

tile_id = tlx.clc_consumer(clc_context, clc_buf, clc_phase)

# clc_phase ^= 1

clc_phase = clc_phase ^ (clc_buf == (NUM_CLC_STAGES - 1))

clc_buf += 1

with tlx.async_task(num_warps=1, num_regs=232): # producer, TMA load

# common code duplicated for each region to avoid SMEM overhead

start_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

num_tiles = num_pid_m * num_pid_n

k_tiles = tl.cdiv(K, BLOCK_SIZE_K)

# end of common code

load_phase = 0 # the current phase of TMA load

# we virtually "flatten" the two layer loop as if we're performing tma loads on

# one big list of data

processed_k_iters = 0

for tile_id in range(start_pid, num_tiles, NUM_SMS):

tile_id = start_pid

clc_phase = 0

clc_buf = 0

while tile_id != -1:

clc_buf = clc_buf % NUM_CLC_STAGES

pid_m, pid_n = _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M)

offs_am = pid_m * BLOCK_SIZE_M

offs_bn = pid_n * BLOCK_SIZE_N

for k in range(0, k_tiles):

# processed_k_iters + k means we use the immediate next buffer slot of tile_id x when we start tile_id x+1

buf = (processed_k_iters + k) % NUM_SMEM_BUFFERS

# wait for previous phase(round) of dot for this buf

tlx.barrier_wait(smem_empty_bars[buf], load_phase ^ 1)

# buffer is now ready to be used again

offs_k = k * BLOCK_SIZE_K

tlx.barrier_expect_bytes(smem_full_bars[buf],

2 * (BLOCK_SIZE_M + BLOCK_SIZE_N) * BLOCK_SIZE_K) # float16

tlx.async_descriptor_load(a_desc, buffers_A[buf], [offs_am, offs_k], smem_full_bars[buf])

tlx.async_descriptor_load(b_desc, buffers_B[buf], [offs_k, offs_bn], smem_full_bars[buf])

# flip phase at the end of a round

load_phase = load_phase ^ (buf == NUM_SMEM_BUFFERS - 1)

processed_k_iters += k_tiles

tile_id = tlx.clc_consumer(clc_context, clc_buf, clc_phase)

# clc_phase ^= 1

clc_phase = clc_phase ^ (clc_buf == (NUM_CLC_STAGES - 1))

clc_buf += 1

Diferenças salvas

Texto original

Abrir arquivo

@triton.jit
def matmul_kernel_tma_ws_blackwell_persistent(a_desc, b_desc, c_desc, M, N, K, BLOCK_SIZE_M: tl.constexpr,
                                              BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,  #
                                              GROUP_SIZE_M: tl.constexpr,  #
                                              NUM_SMEM_BUFFERS: tl.constexpr,  #
                                              NUM_TMEM_BUFFERS: tl.constexpr,  #
                                              NUM_SMS: tl.constexpr,  #
                                              EPILOGUE_SUBTILE: tl.constexpr,  #
                                              ):
    # allocate NUM_SMEM_BUFFERS buffers
    buffers_A = tlx.local_alloc((BLOCK_SIZE_M, BLOCK_SIZE_K), tl.float16, NUM_SMEM_BUFFERS)
    buffers_B = tlx.local_alloc((BLOCK_SIZE_K, BLOCK_SIZE_N), tl.float16, NUM_SMEM_BUFFERS)
    # use multiple TMEM buffers to overlap MMA and epilogue
    tmem_buffers = tlx.local_alloc((BLOCK_SIZE_M, BLOCK_SIZE_N), tl.float32, NUM_TMEM_BUFFERS, tlx.storage_kind.tmem)

# allocate barriers
    smem_empty_bars = tlx.alloc_barriers(num_barriers=NUM_SMEM_BUFFERS, arrive_count=1)
    smem_full_bars = tlx.alloc_barriers(num_barriers=NUM_SMEM_BUFFERS, arrive_count=1)
    tmem_full_bars = tlx.alloc_barriers(num_barriers=NUM_TMEM_BUFFERS, arrive_count=1)
    tmem_empty_bars = tlx.alloc_barriers(num_barriers=NUM_TMEM_BUFFERS, arrive_count=1)

with tlx.async_tasks():
        with tlx.async_task("default"):  # epilogue consumer
            # common code duplicated for each region to avoid SMEM overhead
            start_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
            num_tiles = num_pid_m * num_pid_n
            k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
            # end of common code

tmem_read_phase = 0
            cur_tmem_buf = 0

for tile_id in range(start_pid, num_tiles, NUM_SMS):
                pid_m, pid_n = _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M)
                offs_am = pid_m * BLOCK_SIZE_M
                offs_bn = pid_n * BLOCK_SIZE_N

tlx.barrier_wait(tmem_full_bars[cur_tmem_buf], tmem_read_phase)
                # flip phase at the end of a round of using TMEM barriers
                tmem_read_phase = tmem_read_phase ^ (cur_tmem_buf == NUM_TMEM_BUFFERS - 1)

# load the result from TMEM to registers
                acc_tmem = tmem_buffers[cur_tmem_buf]

if EPILOGUE_SUBTILE:
                    # We load/store the result half by half to reduce SMEM pressure
                    acc_tmem_subslice1 = tlx.subslice(acc_tmem, 0, BLOCK_SIZE_N // 2)
                    result = tlx.local_load(acc_tmem_subslice1)
                    c = result.to(tl.float16)
                    c_desc.store([offs_am, offs_bn], c)

acc_tmem_subslice2 = tlx.subslice(acc_tmem, BLOCK_SIZE_N // 2, BLOCK_SIZE_N // 2)
                    result = tlx.local_load(acc_tmem_subslice2)
                    c = result.to(tl.float16)
                    c_desc.store([offs_am, offs_bn + BLOCK_SIZE_N // 2], c)
                else:
                    result = tlx.local_load(acc_tmem)
                    c = result.to(tl.float16)
                    c_desc.store([offs_am, offs_bn], c)

# done storing this buffer, signal MMA consumer to resume writing to it
                tlx.barrier_arrive(tmem_empty_bars[cur_tmem_buf], 1)

cur_tmem_buf = (cur_tmem_buf + 1) % NUM_TMEM_BUFFERS

with tlx.async_task(num_warps=1, num_regs=232):  # MMA consumer
            # common code duplicated for each region to avoid SMEM overhead
            start_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
            num_tiles = num_pid_m * num_pid_n
            k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
            # end of common code

dot_phase = 0  # the current phase of dot op
            tmem_write_phase = 1  # sync between epilogue consumer and MMA consumer
            cur_tmem_buf = 0

processed_k_iters = 0
            for tile_id in range(start_pid, num_tiles, NUM_SMS):
                pid_m, pid_n = _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M)
                offs_am = pid_m * BLOCK_SIZE_M
                offs_bn = pid_n * BLOCK_SIZE_N

# wait epilogue consumer to be done with the buffer before reusing it
                tlx.barrier_wait(tmem_empty_bars[cur_tmem_buf], tmem_write_phase)
                # flip phase at the end of a round of using TMEM barriers
                tmem_write_phase = tmem_write_phase ^ (cur_tmem_buf == NUM_TMEM_BUFFERS - 1)

# now iterate along K to compute result for the block
                for k in range(0, k_tiles):
                    # processed_k_iters + k means we use the immediate next buffer slot of tile_id x when we start tile_id x+1
                    buf = (processed_k_iters + k) % NUM_SMEM_BUFFERS
                    # wait for current phase(round) of load for this buf
                    tlx.barrier_wait(smem_full_bars[buf], dot_phase)
                    # buffer is now ready with loaded data, tlx.async_dot will signal `mBarrier` when done
                    tlx.async_dot(
                        buffers_A[buf],
                        buffers_B[buf],
                        tmem_buffers[cur_tmem_buf],
                        use_acc=k > 0,
                        mBarriers=[smem_empty_bars[buf]],
                        out_dtype=tl.float32,
                    )
                    # flip phase at the end of a round
                    dot_phase = dot_phase ^ (buf == NUM_SMEM_BUFFERS - 1)

# wait for last mma to complete
                last_buf = (processed_k_iters + k_tiles - 1) % NUM_SMEM_BUFFERS
                # in case phase was flipped, we should use the phase value when dot op was issued
                last_dot_phase = dot_phase ^ (last_buf == NUM_SMEM_BUFFERS - 1)
                tlx.barrier_wait(smem_empty_bars[last_buf], last_dot_phase)

# done filling this buffer, signal epilogue consumer
                tlx.barrier_arrive(tmem_full_bars[cur_tmem_buf], 1)

# possibly enter next iteration (next tile) without waiting for epilogue
                cur_tmem_buf = (cur_tmem_buf + 1) % NUM_TMEM_BUFFERS
                processed_k_iters += k_tiles

with tlx.async_task(num_warps=1, num_regs=232):  # producer, TMA load
            # common code duplicated for each region to avoid SMEM overhead
            start_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
            num_tiles = num_pid_m * num_pid_n
            k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
            # end of common code

load_phase = 0  # the current phase of TMA load
            # we virtually "flatten" the two layer loop as if we're performing tma loads on
            # one big list of data
            processed_k_iters = 0

for k in range(0, k_tiles):
                    # processed_k_iters + k means we use the immediate next buffer slot of tile_id x when we start tile_id x+1
                    buf = (processed_k_iters + k) % NUM_SMEM_BUFFERS
                    # wait for previous phase(round) of dot for this buf
                    tlx.barrier_wait(smem_empty_bars[buf], load_phase ^ 1)
                    # buffer is now ready to be used again
                    offs_k = k * BLOCK_SIZE_K
                    tlx.barrier_expect_bytes(smem_full_bars[buf],
                                             2 * (BLOCK_SIZE_M + BLOCK_SIZE_N) * BLOCK_SIZE_K)  # float16
                    tlx.async_descriptor_load(a_desc, buffers_A[buf], [offs_am, offs_k], smem_full_bars[buf])
                    tlx.async_descriptor_load(b_desc, buffers_B[buf], [offs_k, offs_bn], smem_full_bars[buf])
                    # flip phase at the end of a round
                    load_phase = load_phase ^ (buf == NUM_SMEM_BUFFERS - 1)
                processed_k_iters += k_tiles

Texto alterado

Abrir arquivo

def matmul_kernel_tma_ws_blackwell_clc(a_desc, b_desc, c_desc, M, N, K, BLOCK_SIZE_M: tl.constexpr,
                                       BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,  #
                                       GROUP_SIZE_M: tl.constexpr,  #
                                       NUM_SMEM_BUFFERS: tl.constexpr,  #
                                       NUM_TMEM_BUFFERS: tl.constexpr,  #
                                       NUM_SMS: tl.constexpr,  #
                                       NUM_CLC_STAGES: tl.constexpr,  #
                                       EPILOGUE_SUBTILE: tl.constexpr,  #
                                       ):
    # allocate NUM_SMEM_BUFFERS buffers
    buffers_A = tlx.local_alloc((BLOCK_SIZE_M, BLOCK_SIZE_K), tl.float16, NUM_SMEM_BUFFERS)
    buffers_B = tlx.local_alloc((BLOCK_SIZE_K, BLOCK_SIZE_N), tl.float16, NUM_SMEM_BUFFERS)
    # use multiple TMEM buffers to overlap MMA and epilogue
    tmem_buffers = tlx.local_alloc((BLOCK_SIZE_M, BLOCK_SIZE_N), tl.float32, NUM_TMEM_BUFFERS, tlx.storage_kind.tmem)

clc_context = tlx.clc_create_context(NUM_CLC_STAGES, 3)

with tlx.async_tasks():
        with tlx.async_task("default"):  # epilogue consumer
            # common code duplicated for each region to avoid SMEM overhead
            start_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
            k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
            # end of common code

tmem_read_phase = 0
            cur_tmem_buf = 0

tile_id = start_pid

clc_phase_producer = 1
            clc_phase_consumer = 0
            clc_buf = 0
            while tile_id != -1:
                clc_buf = clc_buf % NUM_CLC_STAGES
                # Debug prints
                # if tlx.thread_id(axis=0) == 0:
                # tl.device_print("Default WG Processing CtaID", tile_id)
                # producer
                tlx.clc_producer(clc_context, clc_buf, clc_phase_producer)
                # clc_phase_producer ^= 1
                clc_phase_producer = clc_phase_producer ^ (clc_buf == (NUM_CLC_STAGES - 1))

pid_m, pid_n = _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M)
                offs_am = pid_m * BLOCK_SIZE_M
                offs_bn = pid_n * BLOCK_SIZE_N

# load the result from TMEM to registers
                acc_tmem = tmem_buffers[cur_tmem_buf]

# done storing this buffer, signal MMA consumer to resume writing to it
                tlx.barrier_arrive(tmem_empty_bars[cur_tmem_buf], 1)

cur_tmem_buf = (cur_tmem_buf + 1) % NUM_TMEM_BUFFERS

tile_id = tlx.clc_consumer(clc_context, clc_buf, clc_phase_consumer)
                # clc_phase_consumer ^= 1
                clc_phase_consumer = clc_phase_consumer ^ (clc_buf == (NUM_CLC_STAGES - 1))
                clc_buf += 1

# Debug-only: verifying that CLC steals workloads successfully
                # if tlx.thread_id(axis=0) == 0:
                # tl.device_print("Extracted CtaID", tile_id)

with tlx.async_task(num_warps=1, num_regs=232):  # MMA consumer
            # common code duplicated for each region to avoid SMEM overhead
            start_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
            k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
            # end of common code

dot_phase = 0  # the current phase of dot op
            tmem_write_phase = 1  # sync between epilogue consumer and MMA consumer
            cur_tmem_buf = 0

processed_k_iters = 0
            tile_id = start_pid
            clc_phase = 0
            clc_buf = 0
            while tile_id != -1:
                clc_buf = clc_buf % NUM_CLC_STAGES
                pid_m, pid_n = _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M)
                offs_am = pid_m * BLOCK_SIZE_M
                offs_bn = pid_n * BLOCK_SIZE_N

# now iterate along K to compute result for the block
                for k in range(0, k_tiles):
                    # processed_k_iters + k means we use the immediate next buffer slot of tile_id x when we start tile_id x+1
                    buf = (processed_k_iters + k) % NUM_SMEM_BUFFERS
                    # wait for current phase(round) of load for this buf
                    tlx.barrier_wait(smem_full_bars[buf], dot_phase)
                    # buffer is now ready with loaded data, tlx.async_dot will signal `mBarrier` when done
                    tlx.async_dot(buffers_A[buf], buffers_B[buf], tmem_buffers[cur_tmem_buf], use_acc=k > 0,
                                  mBarriers=[smem_empty_bars[buf]], out_dtype=tl.float32)
                    # flip phase at the end of a round
                    dot_phase = dot_phase ^ (buf == NUM_SMEM_BUFFERS - 1)

# done filling this buffer, signal epilogue consumer
                tlx.barrier_arrive(tmem_full_bars[cur_tmem_buf], 1)

# possibly enter next iteration (next tile) without waiting for epilogue
                cur_tmem_buf = (cur_tmem_buf + 1) % NUM_TMEM_BUFFERS
                processed_k_iters += k_tiles
                tile_id = tlx.clc_consumer(clc_context, clc_buf, clc_phase)
                # clc_phase ^= 1
                clc_phase = clc_phase ^ (clc_buf == (NUM_CLC_STAGES - 1))
                clc_buf += 1

with tlx.async_task(num_warps=1, num_regs=232):  # producer, TMA load
            # common code duplicated for each region to avoid SMEM overhead
            start_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
            k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
            # end of common code

tile_id = start_pid
            clc_phase = 0
            clc_buf = 0
            while tile_id != -1:
                clc_buf = clc_buf % NUM_CLC_STAGES
                pid_m, pid_n = _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M)
                offs_am = pid_m * BLOCK_SIZE_M
                offs_bn = pid_n * BLOCK_SIZE_N