टेक्स्ट की तुलना करें

दो टेक्स्ट फ़ाइलों के बीच अंतर ढूंढें

रियल-टाइम एडिटर

रिक्त स्थान छिपाएँ

अपरिवर्तित संक्षिप्त करें

लाइन रैप बंद

लेआउट

परिवर्तन हाइलाइट करें

टेक्स्ट शैलियां

सिंटैक्स हाइलाइटिंग

Diffchecker Desktop Diffchecker चलाने का सबसे सुरक्षित तरीका। Diffchecker Desktop ऐप पाएं: आपके diffs कभी आपके कंप्यूटर से बाहर नहीं जाते!Desktop पाएं

Comparing TLX between persistent and CLC

बनाया गया 8 माह पहलेDiff कभी समाप्त नहीं होता

35 हटाए गए

लाइनें
कुल
हटाया गया

अक्षर
कुल
हटाया गया

इस सुविधा का उपयोग जारी रखने के लिए, अपग्रेड करें Diffchecker Pro मूल्य देखें

158 लाइनें

54 जोड़े गए

लाइनें
कुल
जोड़ा गया

अक्षर
कुल
जोड़ा गया

इस सुविधा का उपयोग जारी रखने के लिए, अपग्रेड करें Diffchecker Pro मूल्य देखें

191 लाइनें

@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

सेव किए गए Diffs

ऑरिजनल टेक्स्ट

फ़ाइल खोलें

@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

परिवर्तित टेक्स्ट

फ़ाइल खोलें

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