From 4f54ad5340f60e18d776dbd6abb7a715e1f75caf Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Wed, 25 Mar 2026 13:47:29 +0800 Subject: [PATCH 1/6] Initialize basic changes for Let defined variables --- examples/example_flex_attn_wgmma.py | 647 ++++++++++++++++++ examples/example_flex_attn_wgmma_unchanged.py | 635 +++++++++++++++++ src/transform/auto_schedule.cc | 415 +++++++---- src/transform/auto_schedule/ir_structure.cc | 43 +- src/transform/auto_schedule/ir_structure.h | 132 +++- src/transform/auto_schedule/memory_detector.h | 51 +- src/transform/if_condition_extract.cc | 123 ++++ tilelang/engine/phase.py | 7 +- tilelang/transform/__init__.py | 3 + 9 files changed, 1900 insertions(+), 156 deletions(-) create mode 100644 examples/example_flex_attn_wgmma.py create mode 100644 examples/example_flex_attn_wgmma_unchanged.py create mode 100644 src/transform/if_condition_extract.cc diff --git a/examples/example_flex_attn_wgmma.py b/examples/example_flex_attn_wgmma.py new file mode 100644 index 0000000000..12accb7dd9 --- /dev/null +++ b/examples/example_flex_attn_wgmma.py @@ -0,0 +1,647 @@ +import argparse +import itertools +from typing import Callable + +import tilelang +import tilelang.language as T +import torch + + +@tilelang.jit( + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + # tilelang.PassConfigKey.TL_ENABLE_AUTO_SCHEDULE: True + } +) +def flashattn_fwd( + batch: int, + heads: int, + dim_qk: int, + dim_vo: int, + softmax_scale: float, + mask_fn: Callable[[int, int, int, int, int, int, T.ptr], bool], + block_mask_fn: Callable[[int, int, int, int, int, int, int, int, T.ptr], bool], + block_qo: int = 64, + block_kv: int = 64, + num_stages: int = 2, + thread_num: int = 128, +): + scale = softmax_scale * 1.44269504 # log2(e) + dtype = T.bfloat16 + accum_dtype = T.float32 + + seq_len_qo = T.dynamic('seq_len_qo') + seq_len_kv = T.dynamic('seq_len_kv') + + @T.prim_func + def flash_fwd( + q_global: T.Tensor([batch, seq_len_qo, heads, dim_qk], dtype), + k_global: T.Tensor([batch, seq_len_kv, heads, dim_qk], dtype), + v_global: T.Tensor([batch, seq_len_kv, heads, dim_vo], dtype), + o_global: T.Tensor([batch, seq_len_qo, heads, dim_vo], dtype), + lse_global: T.Tensor([batch, heads, seq_len_qo], accum_dtype), + mask_custom_data: T.ptr, + block_mask_custom_data: T.ptr, + ): + with T.Kernel(T.ceildiv(seq_len_qo, block_qo), heads, batch, threads=thread_num) as ( + bx, + by, + bz, + ): + q_shared = T.alloc_shared([block_qo, dim_qk], dtype) + k_shared = T.alloc_shared([block_kv, dim_qk], dtype) + v_shared = T.alloc_shared([block_kv, dim_vo], dtype) + acc_s = T.alloc_fragment([block_qo, block_kv], accum_dtype) + acc_s_cast = T.alloc_fragment([block_qo, block_kv], dtype) + acc_o = T.alloc_fragment([block_qo, dim_vo], accum_dtype) + scores_max = T.alloc_fragment([block_qo], accum_dtype) + scores_max_prev = T.alloc_fragment([block_qo], accum_dtype) + scores_scale = T.alloc_fragment([block_qo], accum_dtype) + scores_sum = T.alloc_fragment([block_qo], accum_dtype) + scaled_sum_exp = T.alloc_fragment([block_qo], accum_dtype) + + T.copy(q_global[bz, bx * block_qo : (bx + 1) * block_qo, by, :], q_shared) + T.fill(acc_o, 0) + T.fill(scaled_sum_exp, 0) + T.fill(scores_max, -2.**100) + + for k in T.Pipelined(T.ceildiv(seq_len_kv, block_kv), num_stages=num_stages): + cond = block_mask_fn( + bz, + by, + bx * block_qo, + T.min(seq_len_qo, (bx + 1) * block_qo), + k * block_kv, + T.min(seq_len_kv, (k + 1) * block_kv), + seq_len_qo, + seq_len_kv, + block_mask_custom_data, + ) + if cond: + T.copy(k_global[bz, k * block_kv : (k + 1) * block_kv, by, :], k_shared) + for i, j in T.Parallel(block_qo, block_kv): + acc_s[i, j] = T.if_then_else( + mask_fn( + bz, + by, + bx * block_qo + i, + k * block_kv + j, + seq_len_qo, + seq_len_kv, + mask_custom_data, + ) + and k * block_kv + j < seq_len_kv, + 0, + -2.**100, + ) + T.gemm( + q_shared, + k_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.copy(v_global[bz, k * block_kv : (k + 1) * block_kv, by, :], v_shared) + T.copy(scores_max, scores_max_prev) + T.reduce_max(acc_s, scores_max, dim=1, clear=False) + for i in T.Parallel(block_qo): + scores_max[i] = T.max(scores_max[i], scores_max_prev[i]) + for i in T.Parallel(block_qo): + scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale) + for i, j in T.Parallel(block_qo, dim_vo): + acc_o[i, j] *= scores_scale[i] + for i, j in T.Parallel(block_qo, block_kv): + acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale) + T.copy(acc_s, acc_s_cast) + T.gemm(acc_s_cast, v_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) + T.reduce_sum(acc_s, scores_sum, dim=1) + for i in T.Parallel(block_qo): + scaled_sum_exp[i] = scaled_sum_exp[i] * scores_scale[i] + scores_sum[i] + + for i, j in T.Parallel(block_qo, dim_vo): + acc_o[i, j] /= scaled_sum_exp[i] + T.copy(acc_o, o_global[bz, bx * block_qo : (bx + 1) * block_qo, by, :]) + + for i in T.Parallel(block_qo): + scaled_sum_exp[i] = T.log2(scaled_sum_exp[i]) + scores_max[i] * scale + T.copy(scaled_sum_exp, lse_global[bz, by, bx * block_qo : (bx + 1) * block_qo]) + + return flash_fwd + + +@tilelang.jit( + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + # tilelang.PassConfigKey.TL_ENABLE_AUTO_SCHEDULE: True + }, +) +def flashattn_bwd_preprocess(batch, heads, dim): + dtype = T.bfloat16 + accum_dtype = T.float32 + seq_len = T.dynamic('seq_len') + shape = [batch, seq_len, heads, dim] + blk = 32 + + @T.prim_func + def flash_bwd_prep( + o_global: T.Tensor(shape, dtype), + grad_o_global: T.Tensor(shape, dtype), + delta_global: T.Tensor([batch, heads, seq_len], accum_dtype), + ): + with T.Kernel(heads, T.ceildiv(seq_len, blk), batch) as (bx, by, bz): + o = T.alloc_fragment([blk, blk], dtype) + do = T.alloc_fragment([blk, blk], dtype) + acc = T.alloc_fragment([blk, blk], accum_dtype) + delta = T.alloc_fragment([blk], accum_dtype) + T.clear(acc) + for k in range(T.ceildiv(dim, blk)): + T.copy( + o_global[ + bz, + by * blk : (by + 1) * blk, + bx, + k * blk : (k + 1) * blk, + ], + o, + ) + T.copy( + grad_o_global[ + bz, + by * blk : (by + 1) * blk, + bx, + k * blk : (k + 1) * blk, + ], + do, + ) + for i, j in T.Parallel(blk, blk): + acc[i, j] += o[i, j] * do[i, j] + T.reduce_sum(acc, delta, 1) + T.copy(delta, delta_global[bz, bx, by * blk : (by + 1) * blk]) + + return flash_bwd_prep + + +@tilelang.jit( + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + # tilelang.PassConfigKey.TL_ENABLE_AUTO_SCHEDULE: True + } +) +def flashattn_bwd( + batch: int, + heads: int, + dim_qk: int, + dim_vo: int, + softmax_scale: float, + mask_fn: Callable[[int, int, int, int, int, int], bool], + block_mask_fn: Callable[[int, int, int, int, int, int, int, int], bool], + block_qo: int = 32, + block_kv: int = 64, + num_stages: int = 2, + thread_num: int = 128, +): + scale = softmax_scale * 1.44269504 # log2(e) + dtype = T.bfloat16 + accum_dtype = T.float32 + + seq_len_qo = T.dynamic('seq_len_qo') + seq_len_kv = T.dynamic('seq_len_kv') + + @T.prim_func + def flash_bwd( + q_global: T.Tensor([batch, seq_len_qo, heads, dim_qk], dtype), + k_global: T.Tensor([batch, seq_len_kv, heads, dim_qk], dtype), + v_global: T.Tensor([batch, seq_len_kv, heads, dim_vo], dtype), + grad_o_global: T.Tensor([batch, seq_len_qo, heads, dim_vo], dtype), + lse_global: T.Tensor([batch, heads, seq_len_qo], accum_dtype), + delta_global: T.Tensor([batch, heads, seq_len_qo], accum_dtype), + grad_q_global: T.Tensor( + [batch, heads, T.ceildiv(seq_len_qo, block_qo) * block_qo, dim_qk], accum_dtype + ), + grad_k_global: T.Tensor([batch, seq_len_kv, heads, dim_qk], dtype), + grad_v_global: T.Tensor([batch, seq_len_kv, heads, dim_vo], dtype), + mask_custom_data: T.ptr, + block_mask_custom_data: T.ptr, + ): + with T.Kernel(heads, T.ceildiv(seq_len_kv, block_kv), batch, threads=thread_num) as ( + bx, + by, + bz, + ): + T.disable_warp_group_reg_alloc() + K_shared = T.alloc_shared([block_kv, dim_qk], dtype) + dsT_shared = T.alloc_shared([block_kv, block_qo], dtype) + q = T.alloc_shared([block_qo, dim_qk], dtype) + V_shared = T.alloc_shared([block_kv, dim_vo], dtype) + qkT = T.alloc_fragment([block_kv, block_qo], accum_dtype) + dsT = T.alloc_fragment([block_kv, block_qo], accum_dtype) + qkT_cast = T.alloc_shared([block_kv, block_qo], dtype) + lse_shared = T.alloc_shared([block_qo], accum_dtype) + delta = T.alloc_shared([block_qo], accum_dtype) + do = T.alloc_shared([block_qo, dim_vo], dtype) + dv = T.alloc_fragment([block_kv, dim_vo], accum_dtype) + dk = T.alloc_fragment([block_kv, dim_qk], accum_dtype) + dq = T.alloc_fragment([block_qo, dim_qk], accum_dtype) + dv_shared = T.alloc_shared([block_kv, dim_vo], dtype) + dk_shared = T.alloc_shared([block_kv, dim_qk], dtype) + dq_shared = T.alloc_shared([block_qo, dim_qk], accum_dtype) + + T.copy(k_global[bz, by * block_kv : (by + 1) * block_kv, bx, :], K_shared) + T.copy(v_global[bz, by * block_kv : (by + 1) * block_kv, bx, :], V_shared) + T.clear(dv) + T.clear(dk) + for k in T.Pipelined(0, T.ceildiv(seq_len_qo, block_qo), num_stages=num_stages): + if block_mask_fn( + bz, + bx, + k * block_qo, + T.min(seq_len_qo, (k + 1) * block_qo), + by * block_kv, + T.min(seq_len_kv, (by + 1) * block_kv), + seq_len_qo, + seq_len_kv, + block_mask_custom_data, + ): + T.copy(q_global[bz, k * block_qo : (k + 1) * block_qo, bx, :], q) + T.clear(qkT) + T.gemm(K_shared, q, qkT, transpose_B=True) + T.copy(grad_o_global[bz, k * block_qo : (k + 1) * block_qo, bx, :], do) + T.clear(dsT) + T.gemm(V_shared, do, dsT, transpose_B=True) + + T.copy(lse_global[bz, bx, k * block_qo : (k + 1) * block_qo], lse_shared) + for i, j in T.Parallel(block_kv, block_qo): + qkT[i, j] = T.exp2(qkT[i, j] * scale - lse_shared[j]) + # We don't need to handle OOB positions, + # since OOB values won't affect other positions here. + for i, j in T.Parallel(block_kv, block_qo): + qkT[i, j] = T.if_then_else( + mask_fn( + bz, + bx, + k * block_qo + j, + by * block_kv + i, + seq_len_qo, + seq_len_kv, + mask_custom_data, + ), + qkT[i, j], + 0, + ) + T.copy(qkT, qkT_cast) + T.gemm(qkT_cast, do, dv) + + T.copy(delta_global[bz, bx, k * block_qo : (k + 1) * block_qo], delta) + + for i, j in T.Parallel(block_kv, block_qo): + dsT_shared[i, j] = qkT[i, j] * (dsT[i, j] - delta[j]) * softmax_scale + T.gemm(dsT_shared, q, dk) + + T.clear(dq) + T.gemm(dsT_shared, K_shared, dq, transpose_A=True) + T.copy(dq, dq_shared) + T.atomic_add( + grad_q_global[bz, bx, k * block_qo : (k + 1) * block_qo, :], + dq_shared, + ) + T.copy(dv, dv_shared) + T.copy(dk, dk_shared) + T.copy(dv_shared, grad_v_global[bz, by * block_kv : (by + 1) * block_kv, bx, :]) + T.copy(dk_shared, grad_k_global[bz, by * block_kv : (by + 1) * block_kv, bx, :]) + + return flash_bwd + + +class _FlexAttn(torch.autograd.Function): + @staticmethod + def forward( + ctx: torch.autograd.Function, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + softmax_scale: float, + mask_fn: Callable[[int, int], bool], + block_mask_fn: Callable[[int, int, int, int], bool], + mask_custom_data: torch.Tensor, + block_mask_custom_data: torch.Tensor, + ): + batch, seq_len_qo, head, dim_qk = q.shape + _, seq_len_kv, _, dim_vo = v.shape + o = torch.empty((batch, seq_len_qo, head, dim_vo), dtype=q.dtype, device=q.device) + lse = torch.empty((batch, head, seq_len_qo), dtype=torch.float32, device=q.device) + mod = flashattn_fwd( + batch, + head, + dim_qk, + dim_vo, + softmax_scale, + mask_fn, + block_mask_fn, + ) + mod( + q, + k, + v, + o, + lse, + mask_custom_data, + block_mask_custom_data, + ) + ctx.save_for_backward( + q, + k, + v, + o, + lse, + mask_custom_data, + block_mask_custom_data, + ) + ctx.softmax_scale = softmax_scale + ctx.mask_fn = mask_fn + ctx.block_mask_fn = block_mask_fn + return o, lse + + @staticmethod + def backward(ctx: torch.autograd.Function, do: torch.Tensor, dlse: None): + ( + q, + k, + v, + o, + lse, + mask_custom_data, + block_mask_custom_data, + ) = ctx.saved_tensors + batch, seq_len_qo, head, dim_qk = q.shape + _, seq_len_kv, _, dim_vo = v.shape + + do, q, k, v, o = [x.contiguous() for x in (do, q, k, v, o)] + + delta = torch.empty((batch, head, seq_len_qo), dtype=torch.float32, device=q.device) + flashattn_bwd_preprocess(batch, head, dim_vo)(o, do, delta) + + block_qo = 32 + mod = flashattn_bwd( + batch, + head, + dim_qk, + dim_vo, + ctx.softmax_scale, + ctx.mask_fn, + ctx.block_mask_fn, + block_qo=block_qo, + ) + dq = torch.zeros( + (batch, head, (seq_len_qo + block_qo - 1) // block_qo * block_qo, dim_qk), + device=q.device, + dtype=torch.float32, + ) + dk = torch.zeros_like(k) + dv = torch.empty_like(v) + mod( + q, + k, + v, + do, + lse, + delta, + dq, + dk, + dv, + mask_custom_data, + block_mask_custom_data, + ) + dq = dq[:, :, :seq_len_qo, :].transpose(1, 2).bfloat16() # TODO: tilelang cast + return dq, dk, dv, *([None] * 5) + + +flex_attn = _FlexAttn.apply + + +def ref_program( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + softmax_scale: float, + bool_mask: torch.Tensor, +): + score = torch.einsum('bshd,bthd->bhst', q, k) * softmax_scale + score = score.masked_fill(~bool_mask, float('-inf')) + p = torch.softmax(score, dim=-1) + lse = torch.logsumexp(score, dim=-1) * 1.44269504 + o = torch.einsum('bhst,bthd->bshd', p, v) + return o, lse + + +def main( + BATCH: int, + H: int, + N_CTX_QO: int, + D_HEAD_QK: int, + N_CTX_KV: int, + D_HEAD_VO: int, + mode: str, +): + if mode == 'causal': + bool_mask = ( + (N_CTX_QO - torch.arange(N_CTX_QO, device='cuda'))[:, None] + <= (N_CTX_KV - torch.arange(N_CTX_KV, device='cuda'))[None, :] + ).expand(BATCH, H, N_CTX_QO, N_CTX_KV) + + @T.macro + def mask_fn( + b, + h, + q_idx, + k_idx, + seq_len_qo, + seq_len_kv, + mask_custom_data, + ): + return (seq_len_qo - q_idx) <= (seq_len_kv - k_idx) + + @T.macro + def block_mask_fn( + b, + h, + q_start, + q_end, + k_start, + k_end, + seq_len_qo, + seq_len_kv, + block_mask_custom_data, + ): + return (seq_len_qo - q_end - 1) <= (seq_len_kv - k_start) + + mask_custom_data = None + block_mask_custom_data = None + + elif mode == 'full': + bool_mask = torch.ones((BATCH, H, N_CTX_QO, N_CTX_KV), dtype=torch.bool, device='cuda') + + @T.macro + def mask_fn( + b, + h, + q_idx, + k_idx, + seq_len_qo, + seq_len_kv, + mask_custom_data, + ): + return True + + @T.macro + def block_mask_fn( + b, + h, + q_start, + q_end, + k_start, + k_end, + seq_len_qo, + seq_len_kv, + block_mask_custom_data, + ): + return True + + mask_custom_data = None + block_mask_custom_data = None + elif mode == 'random': + bool_mask = torch.rand((BATCH, H, N_CTX_QO, N_CTX_KV), device='cuda') < 0.5 + + @T.macro + def mask_fn( + b, + h, + q_idx, + k_idx, + seq_len_qo, + seq_len_kv, + mask_custom_data, + ): + mask = T.make_tensor( + mask_custom_data, + shape=(BATCH, H, seq_len_qo, seq_len_kv), + dtype=T.bool, + ) + return mask[b, h, q_idx, k_idx] + + @T.macro + def block_mask_fn( + b, + h, + q_start, + q_end, + k_start, + k_end, + seq_len_qo, + seq_len_kv, + block_mask_custom_data, + ): + block_mask = T.make_tensor( + block_mask_custom_data, + shape=(BATCH, H, T.ceildiv(seq_len_qo, 128), T.ceildiv(seq_len_kv, 128)), + dtype=T.bool, + ) + return T.if_then_else( + q_start // 128 == (q_end - 1) // 128 and k_start // 128 == (k_end - 1) // 128, + block_mask[b, h, q_start // 128, k_start // 128], + True, + ) + + mask_custom_data = bool_mask + block_mask_custom_data = ( + bool_mask.unflatten(-2, (-1, 128)).unflatten(-1, (-1, 128)).any(dim=(-1, -3)) + ) + + Q = ( + torch.empty(BATCH, N_CTX_QO, H, D_HEAD_QK, dtype=torch.bfloat16, device='cuda') + .normal_() + .requires_grad_() + ) + K = ( + torch.empty(BATCH, N_CTX_KV, H, D_HEAD_QK, dtype=torch.bfloat16, device='cuda') + .normal_() + .requires_grad_() + ) + V = ( + torch.empty(BATCH, N_CTX_KV, H, D_HEAD_VO, dtype=torch.bfloat16, device='cuda') + .normal_() + .requires_grad_() + ) + dO = ( + torch.empty(BATCH, N_CTX_QO, H, D_HEAD_VO, dtype=torch.bfloat16, device='cuda') + .normal_() + .requires_grad_() + ) + O, lse = flex_attn( + Q, + K, + V, + D_HEAD_QK**-0.5, + mask_fn, + block_mask_fn, + mask_custom_data, + block_mask_custom_data, + ) + O.backward(dO) + dQ, Q.grad = Q.grad.clone(), None + dK, K.grad = K.grad.clone(), None + dV, V.grad = V.grad.clone(), None + + Q_ref = Q.float().detach().requires_grad_() + K_ref = K.float().detach().requires_grad_() + V_ref = V.float().detach().requires_grad_() + O_ref, lse_ref = ref_program(Q_ref, K_ref, V_ref, D_HEAD_QK**-0.5, bool_mask) + O_ref.backward(dO) + dQ_ref, Q_ref.grad = Q_ref.grad.clone(), None + dK_ref, K_ref.grad = K_ref.grad.clone(), None + dV_ref, V_ref.grad = V_ref.grad.clone(), None + + torch.testing.assert_close(lse, lse_ref, rtol=3e-6, atol=1e-5) + torch.testing.assert_close(O.float(), O_ref, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(dV.float(), dV_ref, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(dK.float(), dK_ref, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(dQ.float(), dQ_ref, rtol=1e-2, atol=1e-2) + print('All checks passed.✅') + + def run(): + Q.grad = K.grad = V.grad = None + flex_attn( + Q, + K, + V, + D_HEAD_QK**-0.5, + mask_fn, + block_mask_fn, + mask_custom_data, + block_mask_custom_data, + )[0].backward(dO) + + with torch.profiler.profile() as prof: + for _ in range(5): + run() + print(prof.key_averages().table(sort_by='cuda_time_total', row_limit=10)) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--batch', type=int, default=4, help='Batch size') + parser.add_argument('--h', type=int, default=16, help='Number of heads') + parser.add_argument('--n_ctx_qo', type=int, default=4096, help='Context size') + parser.add_argument('--d_head_qk', type=int, default=192, help='Head dimension') + parser.add_argument('--n_ctx_kv', type=int, default=8192, help='Context size') + parser.add_argument('--d_head_vo', type=int, default=128, help='Head dimension') + parser.add_argument('--mask_mode', type=str, default='random', help='Causal flag') + args = parser.parse_args() + main( + args.batch, + args.h, + args.n_ctx_qo, + args.d_head_qk, + args.n_ctx_kv, + args.d_head_vo, + args.mask_mode, + ) diff --git a/examples/example_flex_attn_wgmma_unchanged.py b/examples/example_flex_attn_wgmma_unchanged.py new file mode 100644 index 0000000000..fe4fdae8ad --- /dev/null +++ b/examples/example_flex_attn_wgmma_unchanged.py @@ -0,0 +1,635 @@ +import argparse +import itertools +from typing import Callable + +import tilelang +import tilelang.language as T +import torch + + +@tilelang.jit(pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True}) +def flashattn_fwd( + batch: int, + heads: int, + dim_qk: int, + dim_vo: int, + softmax_scale: float, + mask_fn: Callable[[int, int, int, int, int, int, T.ptr], bool], + block_mask_fn: Callable[[int, int, int, int, int, int, int, int, T.ptr], bool], + block_qo: int = 64, + block_kv: int = 64, + num_stages: int = 2, + thread_num: int = 128, +): + scale = softmax_scale * 1.44269504 # log2(e) + dtype = T.bfloat16 + accum_dtype = T.float32 + + seq_len_qo = T.dynamic('seq_len_qo') + seq_len_kv = T.dynamic('seq_len_kv') + + @T.prim_func + def flash_fwd( + q_global: T.Tensor([batch, seq_len_qo, heads, dim_qk], dtype), + k_global: T.Tensor([batch, seq_len_kv, heads, dim_qk], dtype), + v_global: T.Tensor([batch, seq_len_kv, heads, dim_vo], dtype), + o_global: T.Tensor([batch, seq_len_qo, heads, dim_vo], dtype), + lse_global: T.Tensor([batch, heads, seq_len_qo], accum_dtype), + mask_custom_data: T.ptr, + block_mask_custom_data: T.ptr, + ): + with T.Kernel(T.ceildiv(seq_len_qo, block_qo), heads, batch, threads=thread_num) as ( + bx, + by, + bz, + ): + q_shared = T.alloc_shared([block_qo, dim_qk], dtype) + k_shared = T.alloc_shared([block_kv, dim_qk], dtype) + v_shared = T.alloc_shared([block_kv, dim_vo], dtype) + acc_s = T.alloc_fragment([block_qo, block_kv], accum_dtype) + acc_s_cast = T.alloc_fragment([block_qo, block_kv], dtype) + acc_o = T.alloc_fragment([block_qo, dim_vo], accum_dtype) + scores_max = T.alloc_fragment([block_qo], accum_dtype) + scores_max_prev = T.alloc_fragment([block_qo], accum_dtype) + scores_scale = T.alloc_fragment([block_qo], accum_dtype) + scores_sum = T.alloc_fragment([block_qo], accum_dtype) + scaled_sum_exp = T.alloc_fragment([block_qo], accum_dtype) + + T.copy(q_global[bz, bx * block_qo : (bx + 1) * block_qo, by, :], q_shared) + T.fill(acc_o, 0) + T.fill(scaled_sum_exp, 0) + T.fill(scores_max, -2.**100) + + for k in T.Pipelined(T.ceildiv(seq_len_kv, block_kv), num_stages=num_stages): + if block_mask_fn( + bz, + by, + bx * block_qo, + T.min(seq_len_qo, (bx + 1) * block_qo), + k * block_kv, + T.min(seq_len_kv, (k + 1) * block_kv), + seq_len_qo, + seq_len_kv, + block_mask_custom_data, + ): + T.copy(k_global[bz, k * block_kv : (k + 1) * block_kv, by, :], k_shared) + for i, j in T.Parallel(block_qo, block_kv): + acc_s[i, j] = T.if_then_else( + mask_fn( + bz, + by, + bx * block_qo + i, + k * block_kv + j, + seq_len_qo, + seq_len_kv, + mask_custom_data, + ) + and k * block_kv + j < seq_len_kv, + 0, + -2.**100, + ) + T.gemm( + q_shared, + k_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.copy(v_global[bz, k * block_kv : (k + 1) * block_kv, by, :], v_shared) + T.copy(scores_max, scores_max_prev) + T.reduce_max(acc_s, scores_max, dim=1, clear=False) + for i in T.Parallel(block_qo): + scores_max[i] = T.max(scores_max[i], scores_max_prev[i]) + for i in T.Parallel(block_qo): + scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale) + for i, j in T.Parallel(block_qo, dim_vo): + acc_o[i, j] *= scores_scale[i] + for i, j in T.Parallel(block_qo, block_kv): + acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale) + T.copy(acc_s, acc_s_cast) + T.gemm(acc_s_cast, v_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) + T.reduce_sum(acc_s, scores_sum, dim=1) + for i in T.Parallel(block_qo): + scaled_sum_exp[i] = scaled_sum_exp[i] * scores_scale[i] + scores_sum[i] + + for i, j in T.Parallel(block_qo, dim_vo): + acc_o[i, j] /= scaled_sum_exp[i] + T.copy(acc_o, o_global[bz, bx * block_qo : (bx + 1) * block_qo, by, :]) + + for i in T.Parallel(block_qo): + scaled_sum_exp[i] = T.log2(scaled_sum_exp[i]) + scores_max[i] * scale + T.copy(scaled_sum_exp, lse_global[bz, by, bx * block_qo : (bx + 1) * block_qo]) + + return flash_fwd + + +@tilelang.jit( + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + }, +) +def flashattn_bwd_preprocess(batch, heads, dim): + dtype = T.bfloat16 + accum_dtype = T.float32 + seq_len = T.dynamic('seq_len') + shape = [batch, seq_len, heads, dim] + blk = 32 + + @T.prim_func + def flash_bwd_prep( + o_global: T.Tensor(shape, dtype), + grad_o_global: T.Tensor(shape, dtype), + delta_global: T.Tensor([batch, heads, seq_len], accum_dtype), + ): + with T.Kernel(heads, T.ceildiv(seq_len, blk), batch) as (bx, by, bz): + o = T.alloc_fragment([blk, blk], dtype) + do = T.alloc_fragment([blk, blk], dtype) + acc = T.alloc_fragment([blk, blk], accum_dtype) + delta = T.alloc_fragment([blk], accum_dtype) + T.clear(acc) + for k in range(T.ceildiv(dim, blk)): + T.copy( + o_global[ + bz, + by * blk : (by + 1) * blk, + bx, + k * blk : (k + 1) * blk, + ], + o, + ) + T.copy( + grad_o_global[ + bz, + by * blk : (by + 1) * blk, + bx, + k * blk : (k + 1) * blk, + ], + do, + ) + for i, j in T.Parallel(blk, blk): + acc[i, j] += o[i, j] * do[i, j] + T.reduce_sum(acc, delta, 1) + T.copy(delta, delta_global[bz, bx, by * blk : (by + 1) * blk]) + + return flash_bwd_prep + + +@tilelang.jit(pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True}) +def flashattn_bwd( + batch: int, + heads: int, + dim_qk: int, + dim_vo: int, + softmax_scale: float, + mask_fn: Callable[[int, int, int, int, int, int], bool], + block_mask_fn: Callable[[int, int, int, int, int, int, int, int], bool], + block_qo: int = 32, + block_kv: int = 64, + num_stages: int = 2, + thread_num: int = 128, +): + scale = softmax_scale * 1.44269504 # log2(e) + dtype = T.bfloat16 + accum_dtype = T.float32 + + seq_len_qo = T.dynamic('seq_len_qo') + seq_len_kv = T.dynamic('seq_len_kv') + + @T.prim_func + def flash_bwd( + q_global: T.Tensor([batch, seq_len_qo, heads, dim_qk], dtype), + k_global: T.Tensor([batch, seq_len_kv, heads, dim_qk], dtype), + v_global: T.Tensor([batch, seq_len_kv, heads, dim_vo], dtype), + grad_o_global: T.Tensor([batch, seq_len_qo, heads, dim_vo], dtype), + lse_global: T.Tensor([batch, heads, seq_len_qo], accum_dtype), + delta_global: T.Tensor([batch, heads, seq_len_qo], accum_dtype), + grad_q_global: T.Tensor( + [batch, heads, T.ceildiv(seq_len_qo, block_qo) * block_qo, dim_qk], accum_dtype + ), + grad_k_global: T.Tensor([batch, seq_len_kv, heads, dim_qk], dtype), + grad_v_global: T.Tensor([batch, seq_len_kv, heads, dim_vo], dtype), + mask_custom_data: T.ptr, + block_mask_custom_data: T.ptr, + ): + with T.Kernel(heads, T.ceildiv(seq_len_kv, block_kv), batch, threads=thread_num) as ( + bx, + by, + bz, + ): + T.disable_warp_group_reg_alloc() + K_shared = T.alloc_shared([block_kv, dim_qk], dtype) + dsT_shared = T.alloc_shared([block_kv, block_qo], dtype) + q = T.alloc_shared([block_qo, dim_qk], dtype) + V_shared = T.alloc_shared([block_kv, dim_vo], dtype) + qkT = T.alloc_fragment([block_kv, block_qo], accum_dtype) + dsT = T.alloc_fragment([block_kv, block_qo], accum_dtype) + qkT_cast = T.alloc_shared([block_kv, block_qo], dtype) + lse_shared = T.alloc_shared([block_qo], accum_dtype) + delta = T.alloc_shared([block_qo], accum_dtype) + do = T.alloc_shared([block_qo, dim_vo], dtype) + dv = T.alloc_fragment([block_kv, dim_vo], accum_dtype) + dk = T.alloc_fragment([block_kv, dim_qk], accum_dtype) + dq = T.alloc_fragment([block_qo, dim_qk], accum_dtype) + dv_shared = T.alloc_shared([block_kv, dim_vo], dtype) + dk_shared = T.alloc_shared([block_kv, dim_qk], dtype) + dq_shared = T.alloc_shared([block_qo, dim_qk], accum_dtype) + + T.copy(k_global[bz, by * block_kv : (by + 1) * block_kv, bx, :], K_shared) + T.copy(v_global[bz, by * block_kv : (by + 1) * block_kv, bx, :], V_shared) + T.clear(dv) + T.clear(dk) + for k in T.Pipelined(0, T.ceildiv(seq_len_qo, block_qo), num_stages=num_stages): + if block_mask_fn( + bz, + bx, + k * block_qo, + T.min(seq_len_qo, (k + 1) * block_qo), + by * block_kv, + T.min(seq_len_kv, (by + 1) * block_kv), + seq_len_qo, + seq_len_kv, + block_mask_custom_data, + ): + T.copy(q_global[bz, k * block_qo : (k + 1) * block_qo, bx, :], q) + T.clear(qkT) + T.gemm(K_shared, q, qkT, transpose_B=True) + T.copy(grad_o_global[bz, k * block_qo : (k + 1) * block_qo, bx, :], do) + T.clear(dsT) + T.gemm(V_shared, do, dsT, transpose_B=True) + + T.copy(lse_global[bz, bx, k * block_qo : (k + 1) * block_qo], lse_shared) + for i, j in T.Parallel(block_kv, block_qo): + qkT[i, j] = T.exp2(qkT[i, j] * scale - lse_shared[j]) + # We don't need to handle OOB positions, + # since OOB values won't affect other positions here. + for i, j in T.Parallel(block_kv, block_qo): + qkT[i, j] = T.if_then_else( + mask_fn( + bz, + bx, + k * block_qo + j, + by * block_kv + i, + seq_len_qo, + seq_len_kv, + mask_custom_data, + ), + qkT[i, j], + 0, + ) + T.copy(qkT, qkT_cast) + T.gemm(qkT_cast, do, dv) + + T.copy(delta_global[bz, bx, k * block_qo : (k + 1) * block_qo], delta) + + for i, j in T.Parallel(block_kv, block_qo): + dsT_shared[i, j] = qkT[i, j] * (dsT[i, j] - delta[j]) * softmax_scale + T.gemm(dsT_shared, q, dk) + + T.clear(dq) + T.gemm(dsT_shared, K_shared, dq, transpose_A=True) + T.copy(dq, dq_shared) + T.atomic_add( + grad_q_global[bz, bx, k * block_qo : (k + 1) * block_qo, :], + dq_shared, + ) + T.copy(dv, dv_shared) + T.copy(dk, dk_shared) + T.copy(dv_shared, grad_v_global[bz, by * block_kv : (by + 1) * block_kv, bx, :]) + T.copy(dk_shared, grad_k_global[bz, by * block_kv : (by + 1) * block_kv, bx, :]) + + return flash_bwd + + +class _FlexAttn(torch.autograd.Function): + @staticmethod + def forward( + ctx: torch.autograd.Function, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + softmax_scale: float, + mask_fn: Callable[[int, int], bool], + block_mask_fn: Callable[[int, int, int, int], bool], + mask_custom_data: torch.Tensor, + block_mask_custom_data: torch.Tensor, + ): + batch, seq_len_qo, head, dim_qk = q.shape + _, seq_len_kv, _, dim_vo = v.shape + o = torch.empty((batch, seq_len_qo, head, dim_vo), dtype=q.dtype, device=q.device) + lse = torch.empty((batch, head, seq_len_qo), dtype=torch.float32, device=q.device) + mod = flashattn_fwd( + batch, + head, + dim_qk, + dim_vo, + softmax_scale, + mask_fn, + block_mask_fn, + ) + mod( + q, + k, + v, + o, + lse, + mask_custom_data, + block_mask_custom_data, + ) + ctx.save_for_backward( + q, + k, + v, + o, + lse, + mask_custom_data, + block_mask_custom_data, + ) + ctx.softmax_scale = softmax_scale + ctx.mask_fn = mask_fn + ctx.block_mask_fn = block_mask_fn + return o, lse + + @staticmethod + def backward(ctx: torch.autograd.Function, do: torch.Tensor, dlse: None): + ( + q, + k, + v, + o, + lse, + mask_custom_data, + block_mask_custom_data, + ) = ctx.saved_tensors + batch, seq_len_qo, head, dim_qk = q.shape + _, seq_len_kv, _, dim_vo = v.shape + + do, q, k, v, o = [x.contiguous() for x in (do, q, k, v, o)] + + delta = torch.empty((batch, head, seq_len_qo), dtype=torch.float32, device=q.device) + flashattn_bwd_preprocess(batch, head, dim_vo)(o, do, delta) + + block_qo = 32 + mod = flashattn_bwd( + batch, + head, + dim_qk, + dim_vo, + ctx.softmax_scale, + ctx.mask_fn, + ctx.block_mask_fn, + block_qo=block_qo, + ) + dq = torch.zeros( + (batch, head, (seq_len_qo + block_qo - 1) // block_qo * block_qo, dim_qk), + device=q.device, + dtype=torch.float32, + ) + dk = torch.zeros_like(k) + dv = torch.empty_like(v) + mod( + q, + k, + v, + do, + lse, + delta, + dq, + dk, + dv, + mask_custom_data, + block_mask_custom_data, + ) + dq = dq[:, :, :seq_len_qo, :].transpose(1, 2).bfloat16() # TODO: tilelang cast + return dq, dk, dv, *([None] * 5) + + +flex_attn = _FlexAttn.apply + + +def ref_program( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + softmax_scale: float, + bool_mask: torch.Tensor, +): + score = torch.einsum('bshd,bthd->bhst', q, k) * softmax_scale + score = score.masked_fill(~bool_mask, float('-inf')) + p = torch.softmax(score, dim=-1) + lse = torch.logsumexp(score, dim=-1) * 1.44269504 + o = torch.einsum('bhst,bthd->bshd', p, v) + return o, lse + + +def main( + BATCH: int, + H: int, + N_CTX_QO: int, + D_HEAD_QK: int, + N_CTX_KV: int, + D_HEAD_VO: int, + mode: str, +): + if mode == 'causal': + bool_mask = ( + (N_CTX_QO - torch.arange(N_CTX_QO, device='cuda'))[:, None] + <= (N_CTX_KV - torch.arange(N_CTX_KV, device='cuda'))[None, :] + ).expand(BATCH, H, N_CTX_QO, N_CTX_KV) + + @T.macro + def mask_fn( + b, + h, + q_idx, + k_idx, + seq_len_qo, + seq_len_kv, + mask_custom_data, + ): + return (seq_len_qo - q_idx) <= (seq_len_kv - k_idx) + + @T.macro + def block_mask_fn( + b, + h, + q_start, + q_end, + k_start, + k_end, + seq_len_qo, + seq_len_kv, + block_mask_custom_data, + ): + return (seq_len_qo - q_end - 1) <= (seq_len_kv - k_start) + + mask_custom_data = None + block_mask_custom_data = None + + elif mode == 'full': + bool_mask = torch.ones((BATCH, H, N_CTX_QO, N_CTX_KV), dtype=torch.bool, device='cuda') + + @T.macro + def mask_fn( + b, + h, + q_idx, + k_idx, + seq_len_qo, + seq_len_kv, + mask_custom_data, + ): + return True + + @T.macro + def block_mask_fn( + b, + h, + q_start, + q_end, + k_start, + k_end, + seq_len_qo, + seq_len_kv, + block_mask_custom_data, + ): + return True + + mask_custom_data = None + block_mask_custom_data = None + elif mode == 'random': + bool_mask = torch.rand((BATCH, H, N_CTX_QO, N_CTX_KV), device='cuda') < 0.5 + + @T.macro + def mask_fn( + b, + h, + q_idx, + k_idx, + seq_len_qo, + seq_len_kv, + mask_custom_data, + ): + mask = T.make_tensor( + mask_custom_data, + shape=(BATCH, H, seq_len_qo, seq_len_kv), + dtype=T.bool, + ) + return mask[b, h, q_idx, k_idx] + + @T.macro + def block_mask_fn( + b, + h, + q_start, + q_end, + k_start, + k_end, + seq_len_qo, + seq_len_kv, + block_mask_custom_data, + ): + block_mask = T.make_tensor( + block_mask_custom_data, + shape=(BATCH, H, T.ceildiv(seq_len_qo, 128), T.ceildiv(seq_len_kv, 128)), + dtype=T.bool, + ) + return T.if_then_else( + q_start // 128 == (q_end - 1) // 128 and k_start // 128 == (k_end - 1) // 128, + block_mask[b, h, q_start // 128, k_start // 128], + True, + ) + + mask_custom_data = bool_mask + block_mask_custom_data = ( + bool_mask.unflatten(-2, (-1, 128)).unflatten(-1, (-1, 128)).any(dim=(-1, -3)) + ) + + Q = ( + torch.empty(BATCH, N_CTX_QO, H, D_HEAD_QK, dtype=torch.bfloat16, device='cuda') + .normal_() + .requires_grad_() + ) + K = ( + torch.empty(BATCH, N_CTX_KV, H, D_HEAD_QK, dtype=torch.bfloat16, device='cuda') + .normal_() + .requires_grad_() + ) + V = ( + torch.empty(BATCH, N_CTX_KV, H, D_HEAD_VO, dtype=torch.bfloat16, device='cuda') + .normal_() + .requires_grad_() + ) + dO = ( + torch.empty(BATCH, N_CTX_QO, H, D_HEAD_VO, dtype=torch.bfloat16, device='cuda') + .normal_() + .requires_grad_() + ) + O, lse = flex_attn( + Q, + K, + V, + D_HEAD_QK**-0.5, + mask_fn, + block_mask_fn, + mask_custom_data, + block_mask_custom_data, + ) + O.backward(dO) + dQ, Q.grad = Q.grad.clone(), None + dK, K.grad = K.grad.clone(), None + dV, V.grad = V.grad.clone(), None + + Q_ref = Q.float().detach().requires_grad_() + K_ref = K.float().detach().requires_grad_() + V_ref = V.float().detach().requires_grad_() + O_ref, lse_ref = ref_program(Q_ref, K_ref, V_ref, D_HEAD_QK**-0.5, bool_mask) + O_ref.backward(dO) + dQ_ref, Q_ref.grad = Q_ref.grad.clone(), None + dK_ref, K_ref.grad = K_ref.grad.clone(), None + dV_ref, V_ref.grad = V_ref.grad.clone(), None + + torch.testing.assert_close(lse, lse_ref, rtol=3e-6, atol=1e-5) + torch.testing.assert_close(O.float(), O_ref, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(dV.float(), dV_ref, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(dK.float(), dK_ref, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(dQ.float(), dQ_ref, rtol=1e-2, atol=1e-2) + print('All checks passed.✅') + + def run(): + Q.grad = K.grad = V.grad = None + flex_attn( + Q, + K, + V, + D_HEAD_QK**-0.5, + mask_fn, + block_mask_fn, + mask_custom_data, + block_mask_custom_data, + )[0].backward(dO) + + with torch.profiler.profile() as prof: + for _ in range(5): + run() + print(prof.key_averages().table(sort_by='cuda_time_total', row_limit=10)) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--batch', type=int, default=4, help='Batch size') + parser.add_argument('--h', type=int, default=16, help='Number of heads') + parser.add_argument('--n_ctx_qo', type=int, default=4096, help='Context size') + parser.add_argument('--d_head_qk', type=int, default=192, help='Head dimension') + parser.add_argument('--n_ctx_kv', type=int, default=8192, help='Context size') + parser.add_argument('--d_head_vo', type=int, default=128, help='Head dimension') + parser.add_argument('--mask_mode', type=str, default='random', help='Causal flag') + args = parser.parse_args() + main( + args.batch, + args.h, + args.n_ctx_qo, + args.d_head_qk, + args.n_ctx_kv, + args.d_head_vo, + args.mask_mode, + ) diff --git a/src/transform/auto_schedule.cc b/src/transform/auto_schedule.cc index 225fcc19bd..5f162e2fdb 100644 --- a/src/transform/auto_schedule.cc +++ b/src/transform/auto_schedule.cc @@ -69,6 +69,34 @@ namespace tl { using namespace tir; using ffi::GetRef; +// Extract all sequencial task nodes from the IR structure tree +void GatherTaskNodesSingle(const std::shared_ptr &node, + std::vector> &task_nodes); +void GatherTaskNodes(const std::vector> &nodes, + std::vector> &task_nodes) { + for (const auto &node : nodes) { + if (node->IsTask()) { + task_nodes.emplace_back(node); + } else if (node->IsSequence()) { + auto seq = static_cast(node.get()); + GatherTaskNodes(seq->children, task_nodes); + } else if (node->IsWrapper()) { + auto wrapper = static_cast(node.get()); + if (wrapper->task) task_nodes.emplace_back(wrapper->task); + if (wrapper->child) GatherTaskNodesSingle(wrapper->child, task_nodes); + } else if (node->IsControl()) { + task_nodes.emplace_back(node); + } else { + LOG(FATAL) << "Unknown node type in GatherTaskNodes"; + } + } +} + +void GatherTaskNodesSingle(const std::shared_ptr &node, + std::vector> &task_nodes) { + return GatherTaskNodes({node}, task_nodes); +} + // Forward declaration Stmt ApplyWarpgroupPartitionToIRStructure( IRStructure *root, IterVar thread_var, std::vector &barrier_buffers, @@ -576,6 +604,21 @@ class ScheduleUnitBuilder { // Z3-based scheduler for loops that calls Python implementation via FFI // with distance-aware dependencies void Z3SchedulePythonLoop(ControlNode *ctrl) { + if (ctrl->child == nullptr) { + LOG(WARNING) << "Z3SchedulePythonLoop called on a control node without child"; + return; + } + std::vector> flat_children; + if (!ctrl->child->IsSequence()) { + GatherTaskNodesSingle(ctrl->child, flat_children); + } else { + auto seq_node = static_cast(ctrl->child.get()); + GatherTaskNodes(seq_node->children, flat_children); + } + auto seq_node = std::make_shared(); + seq_node->children = flat_children; + ctrl->child = std::move(seq_node); + auto seq_body = static_cast(ctrl->child.get()); std::vector nodes; nodes.reserve(seq_body->children.size()); @@ -752,7 +795,7 @@ class ScheduleUnitBuilder { node_to_index[nodes[i]] = i; } - std::vector> reordered_children; + std::vector> reordered_children; reordered_children.reserve(sorted_nodes.size()); for (IRStructure *sorted_node : sorted_nodes) { auto it = node_to_index.find(sorted_node); @@ -766,7 +809,7 @@ class ScheduleUnitBuilder { } for (auto &node : seq_body->children) { - auto unit = std::make_unique(); + auto unit = std::make_shared(); unit->stage = stage_map[node.get()]; unit->child = std::move(node); node = std::move(unit); @@ -827,6 +870,11 @@ class ScheduleUnitBuilder { return a->buffer.same_as(b->buffer); } + // Check if two variables are the same + bool SameVar(const Var &a, const Var &b) const { + return a.same_as(b); + } + // Check if two IRStructures have data dependency (excluding read-after-read) bool HasDependency(const IRStructure *a, const IRStructure *b) const { // Check if either node contains loop_break (if it's a TaskNode) @@ -873,6 +921,22 @@ class ScheduleUnitBuilder { return true; } } + for (const auto &write_var_a : a->GetWriteVars()) { + for (const auto &read_var_b : b->GetReadVars()) { + if (SameVar(write_var_a, read_var_b)) + return true; + } + for (const auto &write_var_b : b->GetWriteVars()) { + if (SameVar(write_var_a, write_var_b)) + return true; + } + } + for (const auto &read_var_a : a->GetReadVars()) { + for (const auto &write_var_b : b->GetWriteVars()) { + if (SameVar(read_var_a, write_var_b)) + return true; + } + } return false; } @@ -926,6 +990,22 @@ class ScheduleUnitBuilder { return true; } } + for (const auto &write_var_a : a->GetWriteVars()) { + for (const auto &read_var_b : b->GetReadVars()) { + if (SameVar(write_var_a, read_var_b)) + return true; + } + for (const auto &write_var_b : b->GetWriteVars()) { + if (SameVar(write_var_a, write_var_b)) + return true; + } + } + for (const auto &read_var_a : a->GetReadVars()) { + for (const auto &write_var_b : b->GetWriteVars()) { + if (SameVar(read_var_a, write_var_b)) + return true; + } + } return false; } @@ -973,21 +1053,12 @@ void ScheduleUnitBuilder::ScheduleRecursive(IRStructure *node) { if (!node) return; - if (node->IsTask()) { - // TaskNode: no further scheduling needed - return; - } else if (node->IsSequence()) { - auto seq = static_cast(node); - - // First, recursively schedule all children - for (size_t i = 0; i < seq->children.size(); ++i) { - ScheduleRecursive(seq->children[i].get()); - } - + auto ChildrenScheduleHelper = [&](std::vector> origin_children) + -> std::vector> { // Now collect child nodes for potential scheduling std::vector child_nodes; - child_nodes.reserve(seq->children.size()); - for (const auto &child : seq->children) { + child_nodes.reserve(origin_children.size()); + for (const auto &child : origin_children) { child_nodes.push_back(child.get()); } @@ -1011,7 +1082,7 @@ void ScheduleUnitBuilder::ScheduleRecursive(IRStructure *node) { } // Reorder children according to sorted_nodes - std::vector> reordered_children; + std::vector> reordered_children; reordered_children.reserve(sorted_nodes.size()); for (IRStructure *sorted_node : sorted_nodes) { @@ -1021,41 +1092,73 @@ void ScheduleUnitBuilder::ScheduleRecursive(IRStructure *node) { "children mapping"; } size_t old_idx = it->second; - reordered_children.push_back(std::move(seq->children[old_idx])); + reordered_children.emplace_back(origin_children[old_idx]); } // Move reordered children back - seq->children = std::move(reordered_children); + origin_children = reordered_children; } - for (auto &node : seq->children) { - auto unit = std::make_unique(); + for (auto &node : origin_children) { + auto unit = std::make_shared(); unit->stage = -1; - unit->child = std::move(node); - node = std::move(unit); + unit->child = std::shared_ptr(node); + node = unit; + } + return origin_children; + }; + + if (node->IsTask()) { + // TaskNode: no further scheduling needed + return; + } else if (node->IsSequence()) { + auto seq = static_cast(node); + + // First, recursively schedule all children + std::vector> seq_children, origin_children; + GatherTaskNodes(seq->children, origin_children); + for (auto &child : origin_children) { + ScheduleRecursive(child.get()); } + + seq->children = ChildrenScheduleHelper(origin_children); return; } else if (node->IsControl()) { auto ctrl = static_cast(node); // Now schedule the ControlNode's internal tasks (if any) as a unit // The body should now be a SequenceNode containing the tasks - if (ctrl->child && ctrl->child->IsSequence()) { - auto seq_body = static_cast(ctrl->child.get()); - for (const auto &child : seq_body->children) { - ScheduleRecursive(child.get()); + if (ctrl->child) { + if (ctrl->child->IsSequence()) { + auto seq_body = static_cast(ctrl->child.get()); + std::vector> origin_children; + GatherTaskNodes(seq_body->children, origin_children); + for (auto &child : origin_children) { + ScheduleRecursive(child.get()); + } + Z3SchedulePythonLoop(ctrl); + } else if (ctrl->child->IsWrapper()) { + auto wrapper = static_cast(ctrl->child.get()); + std::vector>origin_children; + GatherTaskNodes({wrapper->child}, origin_children); + for (auto &child : origin_children) { + ScheduleRecursive(child.get()); + } + Z3SchedulePythonLoop(ctrl); + } else { + ScheduleRecursive(ctrl->child.get()); } - Z3SchedulePythonLoop(ctrl); - } else { - ScheduleRecursive(ctrl->child.get()); } return; } else if (node->IsWrapper()) { auto wrapper = static_cast(node); - if (wrapper->child) { - ScheduleRecursive(wrapper->child.get()); - wrapper->SetII(wrapper->child->GetII()); - wrapper->SetLatency(wrapper->child->GetLatency()); - } + std::vector>origin_children; + GatherTaskNodes({wrapper->child}, origin_children); + for (auto &child : origin_children) { + ScheduleRecursive(child.get()); + } + auto seq_node = std::make_shared(); + seq_node->children = ChildrenScheduleHelper(origin_children); + *node = *(seq_node.get()); return; } @@ -1139,7 +1242,7 @@ class TilelangRootBodyReplacer : public StmtMutator { // Visitor to build IRStructure from TIR statements class IRStructureBuilder : public StmtVisitor { public: - std::unique_ptr Build(const Stmt &stmt, int64_t thread_count = 1, + std::shared_ptr Build(const Stmt &stmt, int64_t thread_count = 1, Target target = Target()) { thread_count_ = thread_count; target_ = target; @@ -1149,7 +1252,7 @@ class IRStructureBuilder : public StmtVisitor { << "IRStructureBuilder: root_ is null after visiting statement. " << "This may indicate an unhandled statement type."; // Return an empty TaskNode as fallback - auto task_node = std::make_unique(); + auto task_node = std::make_shared(); task_node->stmts.push_back(stmt); return task_node; } @@ -1158,7 +1261,7 @@ class IRStructureBuilder : public StmtVisitor { protected: void VisitStmt_(const SeqStmtNode *op) override { - auto seq_node = std::make_unique(); + auto seq_node = std::make_shared(); for (size_t i = 0; i < op->seq.size(); i++) { VisitStmt(op->seq[i]); @@ -1173,7 +1276,7 @@ class IRStructureBuilder : public StmtVisitor { // Determine if this is a sequential or parallel for if (op->kind == ForKind::kSerial) { // Sequential For -> ControlNode - auto control_node = std::make_unique(); + auto control_node = std::make_shared(); control_node->control = GetRef(op); // Process the loop body @@ -1183,17 +1286,10 @@ class IRStructureBuilder : public StmtVisitor { } else { } - // Absorb WrapperNode chain into ControlNode.wrappers - while (control_node->child && control_node->child->IsWrapper()) { - auto *wrapper = static_cast(control_node->child.get()); - control_node->wrappers.push_back(wrapper->wrapper); - control_node->child = std::move(wrapper->child); - } - root_ = std::move(control_node); } else { // Parallel For -> TaskNode - auto task_node = std::make_unique(); + auto task_node = std::make_shared(); task_node->stmts.push_back(GetRef(op)); // Analyze the loop body for resource usage @@ -1205,7 +1301,7 @@ class IRStructureBuilder : public StmtVisitor { void VisitStmt_(const EvaluateNode *op) override { // Evaluate statement (usually a Call) -> TaskNode - auto task_node = std::make_unique(); + auto task_node = std::make_shared(); task_node->stmts.push_back(GetRef(op)); // Analyze the expression for resource usage @@ -1216,7 +1312,7 @@ class IRStructureBuilder : public StmtVisitor { void VisitStmt_(const IfThenElseNode *op) override { // If statement -> treat as TaskNode for now (could be refined later) - auto task_node = std::make_unique(); + auto task_node = std::make_shared(); task_node->stmts.push_back(GetRef(op)); AnalyzeMemoryExpr(op->condition, task_node.get()); @@ -1232,8 +1328,12 @@ class IRStructureBuilder : public StmtVisitor { void VisitStmt_(const LetStmtNode *op) override { // Wrapper statement -> WrapperNode - auto wrapper_node = std::make_unique(); + auto wrapper_node = std::make_shared(); wrapper_node->wrapper = GetRef(op); + auto task_node = std::make_shared(); + task_node->stmts.push_back(GetLetDecl(op)); + AnalyzeResourceUsage(GetLetDecl(op), task_node.get()); + wrapper_node->task = std::move(task_node); // Process the wrapperbody VisitStmt(op->body); @@ -1246,8 +1346,12 @@ class IRStructureBuilder : public StmtVisitor { void VisitStmt_(const AttrStmtNode *op) override { // Wrapper statement -> WrapperNode - auto wrapper_node = std::make_unique(); + auto wrapper_node = std::make_shared(); wrapper_node->wrapper = GetRef(op); + auto task_node = std::make_shared(); + task_node->stmts.push_back(GetAttrDecl(op)); + AnalyzeResourceUsage(GetAttrDecl(op), task_node.get()); + wrapper_node->task = std::move(task_node); // Process the wrapperbody VisitStmt(op->body); @@ -1259,7 +1363,7 @@ class IRStructureBuilder : public StmtVisitor { } void VisitStmt_(const WhileNode *op) override { - auto task_node = std::make_unique(); + auto task_node = std::make_shared(); task_node->stmts.push_back(GetRef(op)); // Analyze condition and body for resource usage @@ -1275,14 +1379,14 @@ class IRStructureBuilder : public StmtVisitor { // TilelangRootBodyExtractor If we encounter it here, it means we're // processing the entire function body (not extracted), which should only // happen when there's no tilelang_root block - auto task_node = std::make_unique(); + auto task_node = std::make_shared(); task_node->stmts.push_back(GetRef(op)); AnalyzeResourceUsage(GetRef(op), task_node.get()); root_ = std::move(task_node); } private: - std::unique_ptr root_; + std::shared_ptr root_; int64_t thread_count_ = 1; Target target_; @@ -1307,7 +1411,7 @@ class IRStructureBuilder : public StmtVisitor { std::vector tensor_core_shapes; // GemmInst: the resolved tensor core instruction (single, asserted - // unique) + // shared) GemmInst gemm_inst{GemmInst::kMMA}; bool has_gemm_inst{false}; @@ -1459,6 +1563,8 @@ class IRStructureBuilder : public StmtVisitor { memory_detector.Analyze(stmt); std::vector read_regions = memory_detector.GetReadRegions(); std::vector write_regions = memory_detector.GetWriteRegions(); + std::vector read_vars = memory_detector.GetReadVars(); + std::vector write_vars = memory_detector.GetWriteVars(); // Merge with existing regions (avoid duplicates) for (const auto ®ion : read_regions) { @@ -1469,6 +1575,14 @@ class IRStructureBuilder : public StmtVisitor { task_node->AddWriteRegion(region); } + for (const auto &var : read_vars) { + task_node->AddReadVar(var); + } + + for (const auto &var : write_vars) { + task_node->AddWriteVar(var); + } + // Estimate latency and initiation interval for this task LatencyEstimator latency_estimator; latency_estimator.SetThreadCount(thread_count_); @@ -1498,6 +1612,8 @@ class IRStructureBuilder : public StmtVisitor { } }; +Stmt ReNestLetStmts(const Stmt &stmt); + // The main pass function tvm::transform::Pass AutoSchedule(const bool enable_epi) { using namespace tir::transform; @@ -1580,6 +1696,9 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) { unit_builder.SetEnableWarpPartition(config.enable_warp_partition); bool double_thread = unit_builder.Build(ir_structure.get()); + PrintIRStructure(ir_structure.get()); + LOG(FATAL) << "Early exit for debugging purposes\n"; + if (!config.enable_warpgroup_partition) { Stmt new_body = ConvertIRStructureToStmt(ir_structure.get(), enable_epi); @@ -1588,6 +1707,8 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) { TilelangRootBodyReplacer replacer(new_body); final_body = replacer(func->body); + final_body = ReNestLetStmts(final_body); + // Create a new PrimFunc with the updated body auto new_func = PrimFunc(func->params, final_body, func->ret_type, func->buffer_map, func->attrs); @@ -1595,7 +1716,7 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) { } // Print the modified summary view - // PrintIRStructure(ir_structure.get()); + PrintIRStructure(ir_structure.get()); // Analyze buffer dependencies and insert barriers before warpgroup // partition @@ -1699,6 +1820,10 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) { final_body = RewriteAllocBuffers(final_body, buffer_infos); } + // LOG(INFO) << final_body << std::endl; + + final_body = ReNestLetStmts(final_body); + // Create a new PrimFunc with the updated body auto new_func = PrimFunc(func->params, final_body, func->ret_type, func->buffer_map, func->attrs); @@ -1709,7 +1834,7 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) { } // Helper function to clone IRStructure with warpgroup filter -std::unique_ptr +std::shared_ptr CloneIRStructureWithWarpgroupFilter(IRStructure *node, int warpgroup_id) { if (!node || !node->containWarpgroupId(warpgroup_id)) return nullptr; @@ -1719,7 +1844,7 @@ CloneIRStructureWithWarpgroupFilter(IRStructure *node, int warpgroup_id) { return task->Clone(); } else if (node->IsSequence()) { auto seq = static_cast(node); - auto new_seq = std::make_unique(); + auto new_seq = std::make_shared(); for (const auto &child : seq->children) { auto new_child = CloneIRStructureWithWarpgroupFilter(child.get(), warpgroup_id); @@ -1730,23 +1855,22 @@ CloneIRStructureWithWarpgroupFilter(IRStructure *node, int warpgroup_id) { return new_seq; } else if (node->IsControl()) { auto ctrl = static_cast(node); - auto new_ctrl = std::make_unique(); + auto new_ctrl = std::make_shared(); new_ctrl->control = ctrl->control; new_ctrl->SetPromote(ctrl->hasPromote()); - new_ctrl->wrappers = ctrl->wrappers; new_ctrl->child = CloneIRStructureWithWarpgroupFilter(ctrl->child.get(), warpgroup_id); return new_ctrl; } else if (node->IsWrapper()) { auto wrapper = static_cast(node); - auto new_wrapper = std::make_unique(); + auto new_wrapper = std::make_shared(); new_wrapper->wrapper = wrapper->wrapper; new_wrapper->child = CloneIRStructureWithWarpgroupFilter(wrapper->child.get(), warpgroup_id); return new_wrapper; } else if (node->IsScheduleUnit()) { auto unit = static_cast(node); - auto new_unit = std::make_unique(); + auto new_unit = std::make_shared(); new_unit->before[warpgroup_id] = unit->before[warpgroup_id]; new_unit->after[warpgroup_id] = unit->after[warpgroup_id]; new_unit->stage = unit->stage; @@ -1869,16 +1993,6 @@ Stmt ConvertIRStructureToStmt(IRStructure *root, const bool outer_enable_epi) { LOG(FATAL); } Stmt body = SeqStmt::Flatten(stmts); - // Re-wrap body with ControlNode wrappers (innermost first = reverse - // order) - for (auto it = ctrl->wrappers.rbegin(); it != ctrl->wrappers.rend(); - ++it) { - if (const auto *let = (*it).as()) { - body = LetStmt(let->var, let->value, body); - } else if (const auto *attr = (*it).as()) { - body = AttrStmt(attr->node, attr->attr_key, attr->value, body); - } - } // Filter out "num_stages" annotation Map filtered_annotations = ctrl->control->annotations; filtered_annotations.erase("num_stages"); @@ -1984,17 +2098,6 @@ Stmt ConvertIRStructureToStmt(IRStructure *root, const bool outer_enable_epi) { steady.push_back(Substitute(stmt, substitution)); } Stmt new_body = SeqStmt::Flatten(steady); - // Re-wrap body with ControlNode wrappers (innermost first = reverse - // order) - for (auto it = ctrl->wrappers.rbegin(); it != ctrl->wrappers.rend(); - ++it) { - if (const auto *let = (*it).as()) { - new_body = LetStmt(let->var, let->value, new_body); - } else if (const auto *attr = (*it).as()) { - new_body = - AttrStmt(attr->node, attr->attr_key, attr->value, new_body); - } - } auto new_var = loop_var.copy_with_suffix(""); // Filter out "num_stages" annotation Map filtered_annotations = ctrl->control->annotations; @@ -2198,16 +2301,6 @@ Stmt ApplyWarpgroupPartitionToIRStructure( LOG(FATAL); } Stmt body = SeqStmt::Flatten(stmts); - // Re-wrap body with ControlNode wrappers (innermost first = reverse - // order) - for (auto it = ctrl->wrappers.rbegin(); it != ctrl->wrappers.rend(); - ++it) { - if (const auto *let = (*it).as()) { - body = LetStmt(let->var, let->value, body); - } else if (const auto *attr = (*it).as()) { - body = AttrStmt(attr->node, attr->attr_key, attr->value, body); - } - } // Filter out "num_stages" annotation Map filtered_annotations = ctrl->control->annotations; filtered_annotations.erase("num_stages"); @@ -2313,17 +2406,6 @@ Stmt ApplyWarpgroupPartitionToIRStructure( steady.push_back(Substitute(stmt, substitution)); } Stmt new_body = SeqStmt::Flatten(steady); - // Re-wrap body with ControlNode wrappers (innermost first = reverse - // order) - for (auto it = ctrl->wrappers.rbegin(); it != ctrl->wrappers.rend(); - ++it) { - if (const auto *let = (*it).as()) { - new_body = LetStmt(let->var, let->value, new_body); - } else if (const auto *attr = (*it).as()) { - new_body = - AttrStmt(attr->node, attr->attr_key, attr->value, new_body); - } - } auto new_var = loop_var.copy_with_suffix(""); // Filter out "num_stages" annotation Map filtered_annotations = ctrl->control->annotations; @@ -2396,11 +2478,11 @@ Stmt ApplyWarpgroupPartitionToIRStructure( // Helper function to clone IRStructure filtering tasks with warpgroup_id == // -1 (neutral tasks) - std::function(IRStructure *)> + std::function(IRStructure *)> clone_neutral_filter; clone_neutral_filter = [&clone_neutral_filter]( - IRStructure *node) -> std::unique_ptr { + IRStructure *node) -> std::shared_ptr { if (!node) return nullptr; @@ -2409,19 +2491,19 @@ Stmt ApplyWarpgroupPartitionToIRStructure( if (task->GetWarpgroupId() == -1) { return task->Clone(); } else { - auto new_task = std::make_unique(); + auto new_task = std::make_shared(); // Empty statements return new_task; } } else if (node->IsSequence()) { auto seq = static_cast(node); - auto new_seq = std::make_unique(); + auto new_seq = std::make_shared(); for (const auto &child : seq->children) { if (child) { auto node = static_cast(child.get()); auto new_node = clone_neutral_filter(node->child.get()); if (new_node) { - auto new_unit = std::make_unique(); + auto new_unit = std::make_shared(); new_unit->child = std::move(new_node); new_seq->children.push_back(std::move(new_unit)); } @@ -2430,7 +2512,7 @@ Stmt ApplyWarpgroupPartitionToIRStructure( return new_seq; } else if (node->IsWrapper()) { auto wrapper = static_cast(node); - auto new_wrapper = std::make_unique(); + auto new_wrapper = std::make_shared(); new_wrapper->child = clone_neutral_filter(wrapper->child.get()); if (new_wrapper->child) { return new_wrapper; @@ -2453,13 +2535,13 @@ Stmt ApplyWarpgroupPartitionToIRStructure( return false; }; - std::function( + std::function( IRStructure *, const std::function &, int)> clone_neutral_filter_with_top_level; clone_neutral_filter_with_top_level = [&clone_neutral_filter_with_top_level, &clone_neutral_filter]( IRStructure *node, const std::function &include_top_level, - int top_level_index) -> std::unique_ptr { + int top_level_index) -> std::shared_ptr { if (!node) return nullptr; @@ -2467,13 +2549,13 @@ Stmt ApplyWarpgroupPartitionToIRStructure( if (include_top_level(top_level_index)) { return clone_neutral_filter(node); } else { - auto new_task = std::make_unique(); + auto new_task = std::make_shared(); // Empty statements return new_task; } } else if (node->IsSequence()) { auto seq = static_cast(node); - auto new_seq = std::make_unique(); + auto new_seq = std::make_shared(); int child_index = 0; for (const auto &child : seq->children) { if (child) { @@ -2484,7 +2566,7 @@ Stmt ApplyWarpgroupPartitionToIRStructure( schedule_unit->child.get(), include_top_level, next_top_level_index); if (new_node) { - auto new_unit = std::make_unique(); + auto new_unit = std::make_shared(); new_unit->child = std::move(new_node); new_seq->children.push_back(std::move(new_unit)); } @@ -2494,7 +2576,7 @@ Stmt ApplyWarpgroupPartitionToIRStructure( return new_seq; } else if (node->IsWrapper()) { auto wrapper = static_cast(node); - auto new_wrapper = std::make_unique(); + auto new_wrapper = std::make_shared(); new_wrapper->child = clone_neutral_filter_with_top_level( wrapper->child.get(), include_top_level, top_level_index); if (new_wrapper->child) { @@ -2675,6 +2757,109 @@ Stmt ApplyWarpgroupPartitionToIRStructure( return combined_stmt; } +// Re-write LetStmt to nest them properly +// Example transformation: +// SeqStmt { +// let x = 42 { Evaluate(0) } // standalone, empty body +// let y = x+1 { Evaluate(0) } // standalone, empty body +// compute(x, y) // actual work +// store(result) +// } +// becomes: +// let x = 42 { +// let y = x+1 { +// SeqStmt { +// compute(x, y) +// store(result) +// } +// } +// } +class LetStmtNester : public StmtMutator { +public: + Stmt VisitStmt_(const SeqStmtNode *op) override { + Array stmts; + for (const auto &stmt : op->seq) { + stmts.push_back(this->VisitStmt(stmt)); + } + + Array flat_stmts; + for (const auto &stmt : stmts) { + if (const auto *inner_seq = stmt.as()) { + for (const auto &inner_stmt : inner_seq->seq) { + flat_stmts.push_back(inner_stmt); + } + } else { + flat_stmts.push_back(stmt); + } + } + stmts = flat_stmts; + + for (int i = static_cast(stmts.size()) - 2; i >= 0; -- i) { + if (const auto *let = stmts[i].as()) { + if (IsEmptyBody(let->body)) { + Stmt absorbed_body = CollectRemaining(stmts, i + 1); + stmts = TruncateAndReplace(stmts, i, LetStmt(let->var, let->value, absorbed_body)); + } + } else if (const auto *attr = stmts[i].as()) { + if (IsEmptyBody(attr->body)) { + Stmt absorbed_body = CollectRemaining(stmts, i + 1); + stmts = TruncateAndReplace(stmts, i, + AttrStmt(attr->node, attr->attr_key, attr->value, absorbed_body)); + } + } + } + + if (stmts.empty()) return Evaluate(0); + if (stmts.size() == 1) return stmts[0]; + + return SeqStmt(stmts); + } + +private: + // Check if a statement body is Evaluate(0) — the empty placeholder + static bool IsEmptyBody(const Stmt &stmt) { + if (const auto *eval = stmt.as()) { + if (const auto *imm = eval->value.as()) { + return imm->value == 0; + } + } + return false; + } + + // Collect stmts[start .. end) into a single Stmt + static Stmt CollectRemaining(const Array &stmts, int start) { + int n = static_cast(stmts.size()); + if (start >= n) { + return Evaluate(0); + } + if (start == n - 1) { + return stmts[start]; + } + Array remaining; + for (int j = start; j < n; ++j) { + remaining.push_back(stmts[j]); + } + return SeqStmt(remaining); + } + + // Keep stmts[0..index), replace stmts[index] with new_stmt, + // discard everything after (already absorbed into new_stmt body) + static Array TruncateAndReplace(const Array &stmts, int index, + Stmt new_stmt) { + Array result; + for (int j = 0; j < index; ++j) { + result.push_back(stmts[j]); + } + result.push_back(new_stmt); + return result; + } +}; + +Stmt ReNestLetStmts(const Stmt &stmt) { + LetStmtNester nester; + return nester(stmt); +} + // StmtMutator to rewrite alloc_buffers in Block nodes class AllocBufferRewriter : public StmtMutator { public: diff --git a/src/transform/auto_schedule/ir_structure.cc b/src/transform/auto_schedule/ir_structure.cc index 4ba3ce9269..45531a6e11 100644 --- a/src/transform/auto_schedule/ir_structure.cc +++ b/src/transform/auto_schedule/ir_structure.cc @@ -117,6 +117,17 @@ std::vector SequenceNode::GetReadRegions() const { return deduplicated; } +std::vector SequenceNode::GetReadVars() const { + std::vector all_vars; + for (const auto &child : children) { + if (child) { + auto child_read_vars = child->GetReadVars(); + all_vars.insert(all_vars.end(), child_read_vars.begin(), child_read_vars.end()); + } + } + return all_vars; +} + std::vector SequenceNode::GetWriteRegions() const { std::vector all_write_regions; for (const auto &child : children) { @@ -144,6 +155,17 @@ std::vector SequenceNode::GetWriteRegions() const { return deduplicated; } +std::vector SequenceNode::GetWriteVars() const { + std::vector all_vars; + for (const auto &child : children) { + if (child) { + auto child_write_vars = child->GetWriteVars(); + all_vars.insert(all_vars.end(), child_write_vars.begin(), child_write_vars.end()); + } + } + return all_vars; +} + int64_t SequenceNode::GetLatency() const { return latency_; } int64_t SequenceNode::GetII() const { return ii_; } @@ -198,8 +220,8 @@ void SequenceNode::AddWriteRegion(const BufferRegion ®ion) { } } -std::unique_ptr SequenceNode::Clone() const { - auto new_seq = std::make_unique(); +std::shared_ptr SequenceNode::Clone() const { + auto new_seq = std::make_shared(); new_seq->children.reserve(children.size()); for (const auto &child : children) { if (child) { @@ -214,8 +236,8 @@ std::unique_ptr SequenceNode::Clone() const { return new_seq; } -std::unique_ptr TaskNode::Clone() const { - auto new_task = std::make_unique(); +std::shared_ptr TaskNode::Clone() const { + auto new_task = std::make_shared(); // Copy statements new_task->stmts = stmts; // Copy resource usage flags @@ -293,8 +315,8 @@ bool TaskNode::ContainsLoopBreak() const { return found_loop_break; } -std::unique_ptr ControlNode::Clone() const { - auto new_ctrl = std::make_unique(); +std::shared_ptr ControlNode::Clone() const { + auto new_ctrl = std::make_shared(); // Copy For control (For is a TVM object with reference counting) new_ctrl->control = control; // Clone child if exists @@ -302,15 +324,14 @@ std::unique_ptr ControlNode::Clone() const { new_ctrl->child = child->Clone(); } // Copy latency and II - new_ctrl->wrappers = wrappers; new_ctrl->SetLatency(GetLatency()); new_ctrl->SetII(GetII()); new_ctrl->SetPromote(hasPromote()); return new_ctrl; } -std::unique_ptr WrapperNode::Clone() const { - auto new_wrapper = std::make_unique(); +std::shared_ptr WrapperNode::Clone() const { + auto new_wrapper = std::make_shared(); // Copy var and value (TVM objects with reference counting) new_wrapper->wrapper = wrapper; // Clone child if exists @@ -323,8 +344,8 @@ std::unique_ptr WrapperNode::Clone() const { return new_wrapper; } -std::unique_ptr ScheduleUnit::Clone() const { - auto new_unit = std::make_unique(); +std::shared_ptr ScheduleUnit::Clone() const { + auto new_unit = std::make_shared(); // Copy var and value (TVM objects with reference counting) new_unit->stage = stage; // Clone child if exists diff --git a/src/transform/auto_schedule/ir_structure.h b/src/transform/auto_schedule/ir_structure.h index 2fce54bb78..715add0a89 100644 --- a/src/transform/auto_schedule/ir_structure.h +++ b/src/transform/auto_schedule/ir_structure.h @@ -28,6 +28,7 @@ class IRStructure; class TaskNode; class ControlNode; class SequenceNode; +class WrapperNode; // Structure to store region access information with warpgroup id struct RegionAccessInfo { @@ -65,7 +66,7 @@ class IRStructure { virtual ~IRStructure() = default; virtual Kind GetKind() const = 0; - virtual std::unique_ptr Clone() const = 0; + virtual std::shared_ptr Clone() const = 0; // Helper methods for safe casting bool IsTask() const { return GetKind() == Kind::kTask; } @@ -83,6 +84,10 @@ class IRStructure { virtual std::vector GetReadRegions() const = 0; virtual std::vector GetWriteRegions() const = 0; + // Variable access (used for dependency analysis) + virtual std::vector GetReadVars() const = 0; + virtual std::vector GetWriteVars() const = 0; + // Latency estimation virtual int64_t GetLatency() const = 0; // Estimated latency in cycles virtual int64_t GetII() const = 0; // Initiation interval in cycles @@ -100,6 +105,10 @@ class IRStructure { virtual void AddReadRegion(const BufferRegion ®ion) = 0; virtual void AddWriteRegion(const BufferRegion ®ion) = 0; + // Helper methods to add variables + virtual void AddReadVar(const Var &var) = 0; + virtual void AddWriteVar(const Var &var) = 0; + // Recursive region collection method virtual void CollectRegions( std::vector &result, @@ -142,6 +151,12 @@ class TaskNode : public IRStructure { std::vector GetWriteRegions() const override { return write_regions_; } + std::vector GetReadVars() const override { + return read_vars_; + } + std::vector GetWriteVars() const override { + return write_vars_; + } // Latency estimation int64_t GetLatency() const override { return latency_; } int64_t GetII() const override { return ii_; } @@ -245,7 +260,7 @@ class TaskNode : public IRStructure { } // Clone method - std::unique_ptr Clone() const override; + std::shared_ptr Clone() const override; // Helper methods to add regions (for incremental analysis) void AddReadRegion(const BufferRegion ®ion) override { @@ -270,6 +285,9 @@ class TaskNode : public IRStructure { write_regions_.push_back(region); } + void AddReadVar(const Var &var) override { read_vars_.push_back(var); } + void AddWriteVar(const Var &var) override { write_vars_.push_back(var); } + void CollectRegions( std::vector &result, std::set>> &visited) const override; @@ -289,6 +307,9 @@ class TaskNode : public IRStructure { std::vector read_regions_; std::vector write_regions_; + std::vector read_vars_; + std::vector write_vars_; + // Latency estimation int64_t latency_{0}; // Estimated latency in cycles int64_t ii_{0}; // Initiation interval in cycles @@ -315,11 +336,7 @@ class TaskNode : public IRStructure { class ControlNode : public IRStructure { public: For control; // The For operation - std::unique_ptr child; - - // Wrapper Stmts (LetStmt/AttrStmt) that originally wrapped the For loop body. - // Stored outermost-first. Re-applied in reverse order during IR conversion. - std::vector wrappers; + std::shared_ptr child; Kind GetKind() const override { return Kind::kControl; } @@ -342,6 +359,14 @@ class ControlNode : public IRStructure { return child ? child->GetWriteRegions() : std::vector{}; } + // Variable access (aggregate from child) + std::vector GetReadVars() const override { + return child ? child->GetReadVars() : std::vector{}; + } + std::vector GetWriteVars() const override { + return child ? child->GetWriteVars() : std::vector{}; + } + // Latency estimation (aggregate from child) int64_t GetLatency() const override { return latency_; } int64_t GetII() const override { return ii_; } @@ -380,6 +405,13 @@ class ControlNode : public IRStructure { child->AddWriteRegion(region); } + void AddReadVar(const Var &var) override { + // ignore + } + void AddWriteVar(const Var &var) override { + // ignore + } + void CollectRegions( std::vector &result, std::set>> &visited) const override; @@ -389,7 +421,7 @@ class ControlNode : public IRStructure { void SetPromote(bool promote) { has_promote_ = promote; } // Clone method - std::unique_ptr Clone() const override; + std::shared_ptr Clone() const override; bool containWarpgroupId(int id) const override { return child->containWarpgroupId(id); @@ -407,7 +439,8 @@ class ControlNode : public IRStructure { class WrapperNode : public IRStructure { public: Stmt wrapper; - std::unique_ptr child; + std::shared_ptr child; + std::shared_ptr task; Kind GetKind() const override { return Kind::kWrapper; } @@ -430,6 +463,14 @@ class WrapperNode : public IRStructure { return child ? child->GetWriteRegions() : std::vector{}; } + // Variable access (aggregate from child) + std::vector GetReadVars() const override { + return child ? child->GetReadVars() : std::vector{}; + } + std::vector GetWriteVars() const override { + return child ? child->GetWriteVars() : std::vector{}; + } + // Latency estimation (aggregate from child) int64_t GetLatency() const override { return latency_; } int64_t GetII() const override { return ii_; } @@ -468,12 +509,19 @@ class WrapperNode : public IRStructure { child->AddWriteRegion(region); } + void AddReadVar(const Var &var) override { + // ignore + } + void AddWriteVar(const Var &var) override { + // ignore + } + void CollectRegions( std::vector &result, std::set>> &visited) const override; // Clone method - std::unique_ptr Clone() const override; + std::shared_ptr Clone() const override; bool containWarpgroupId(int id) const override { return child->containWarpgroupId(id); @@ -489,7 +537,7 @@ class ScheduleUnit : public IRStructure { public: int stage; std::vector> before, after; - std::unique_ptr child; + std::shared_ptr child; ScheduleUnit() { for (unsigned idx = 0; idx != 2; ++idx) { @@ -519,6 +567,13 @@ class ScheduleUnit : public IRStructure { return child ? child->GetWriteRegions() : std::vector{}; } + std::vector GetReadVars() const override { + return child ? child->GetReadVars() : std::vector{}; + } + std::vector GetWriteVars() const override { + return child ? child->GetWriteVars() : std::vector{}; + } + // Latency estimation (aggregate from child) int64_t GetLatency() const override { return latency_; } int64_t GetII() const override { return ii_; } @@ -557,6 +612,15 @@ class ScheduleUnit : public IRStructure { child->AddWriteRegion(region); } + void AddReadVar(const Var &var) override { + if (child) + child->AddReadVar(var); + } + void AddWriteVar(const Var &var) override { + if (child) + child->AddWriteVar(var); + } + void CollectRegions( std::vector &result, std::set>> &visited) const override; @@ -570,7 +634,7 @@ class ScheduleUnit : public IRStructure { } // Clone method - std::unique_ptr Clone() const override; + std::shared_ptr Clone() const override; bool containWarpgroupId(int id) const override { return child->containWarpgroupId(id); @@ -585,7 +649,7 @@ class ScheduleUnit : public IRStructure { // Sequence node: contains a vector of child IRStructures class SequenceNode : public IRStructure { public: - std::vector> children; + std::vector> children; Kind GetKind() const override { return Kind::kSequence; } @@ -598,6 +662,9 @@ class SequenceNode : public IRStructure { std::vector GetReadRegions() const override; std::vector GetWriteRegions() const override; + std::vector GetReadVars() const override; + std::vector GetWriteVars() const override; + // Latency estimation (aggregate from all children) int64_t GetLatency() const override; int64_t GetII() const override; @@ -615,12 +682,19 @@ class SequenceNode : public IRStructure { void AddReadRegion(const BufferRegion ®ion) override; void AddWriteRegion(const BufferRegion ®ion) override; + void AddReadVar(const Var &var) override { + // ignore + } + void AddWriteVar(const Var &var) override { + // ignore + } + void CollectRegions( std::vector &result, std::set>> &visited) const override; // Clone method - std::unique_ptr Clone() const override; + std::shared_ptr Clone() const override; bool containWarpgroupId(int id) const override { for (auto &child : children) { @@ -800,13 +874,6 @@ inline void PrintAllStmts(const IRStructure *node, int indent = 0) { LOG(INFO) << indent_str << " For statement:"; LOG(INFO) << indent_str + " " << control->control; - if (!control->wrappers.empty()) { - LOG(INFO) << indent_str << " Wrappers: " << control->wrappers.size(); - for (size_t w = 0; w < control->wrappers.size(); w++) { - LOG(INFO) << indent_str << " Wrapper " << w << ": " - << control->wrappers[w]; - } - } // Recursively print child statements if (control->child) { LOG(INFO) << indent_str << " Loop body:"; @@ -851,6 +918,20 @@ inline void PrintAllStmts(const IRStructure *node, int indent = 0) { } } +inline Stmt GetLetDecl(const LetStmtNode *let_node) { + if (let_node == nullptr) { + return Stmt(); + } + return LetStmt(let_node->var, let_node->value, Evaluate(0)); +} + +inline Stmt GetAttrDecl(const AttrStmtNode *attr_node) { + if (attr_node == nullptr) { + return Stmt(); + } + return AttrStmt(attr_node->node, attr_node->attr_key, attr_node->value, Evaluate(0)); +} + // Original helper function to print IRStructure (kept for backward // compatibility) inline void PrintIRStructure(const IRStructure *node, int indent = 0) { @@ -882,13 +963,6 @@ inline void PrintIRStructure(const IRStructure *node, int indent = 0) { } else if (node->IsControl()) { const ControlNode *control = static_cast(node); LOG(INFO) << indent_str << "ControlNode (For loop):"; - if (!control->wrappers.empty()) { - LOG(INFO) << indent_str << " Wrappers: " << control->wrappers.size(); - for (size_t w = 0; w < control->wrappers.size(); w++) { - LOG(INFO) << indent_str << " Wrapper " << w << ": " - << control->wrappers[w]; - } - } if (control->child) { LOG(INFO) << indent_str << " Child:"; PrintIRStructure(control->child.get(), indent + 2); @@ -905,6 +979,8 @@ inline void PrintIRStructure(const IRStructure *node, int indent = 0) { const WrapperNode *wrapper = static_cast(node); LOG(INFO) << indent_str << "WrapperNode:"; LOG(INFO) << indent_str << " Wrapper: " << wrapper->wrapper; + LOG(INFO) << indent_str << " Task:"; + PrintIRStructure(wrapper->task.get(), indent + 4); if (wrapper->child) { LOG(INFO) << indent_str << " Child:"; PrintIRStructure(wrapper->child.get(), indent + 2); diff --git a/src/transform/auto_schedule/memory_detector.h b/src/transform/auto_schedule/memory_detector.h index c5e5465d81..dbf1136ebc 100644 --- a/src/transform/auto_schedule/memory_detector.h +++ b/src/transform/auto_schedule/memory_detector.h @@ -38,6 +38,9 @@ class MemoryAccessDetector : public StmtExprVisitor { hint_map_.clear(); pending_conditions_.clear(); let_bindings_.clear(); + read_vars_.clear(); + write_vars_.clear(); + let_defined_vars_.clear(); operator()(stmt); } @@ -50,6 +53,9 @@ class MemoryAccessDetector : public StmtExprVisitor { hint_map_.clear(); pending_conditions_.clear(); let_bindings_.clear(); + read_vars_.clear(); + write_vars_.clear(); + let_defined_vars_.clear(); operator()(expr); } @@ -63,6 +69,11 @@ class MemoryAccessDetector : public StmtExprVisitor { return CollectRegions(write_buffers_, write_regions_); } + // Return all variables that are read from + std::vector GetReadVars() const { return read_vars_; } + // Return all variables that are written to + std::vector GetWriteVars() const { return write_vars_; } + private: /*! \brief Iteration range for loop_vars */ std::unordered_map dom_map_; @@ -83,6 +94,13 @@ class MemoryAccessDetector : public StmtExprVisitor { /*! \brief let bindings inside the block */ std::unordered_map let_bindings_; + /*! \brief The set of variables that are read in the current block. */ + std::vector read_vars_; + /*! \brief The set of variables that are written in the current block. */ + std::vector write_vars_; + /*! \brief The set of variables defined by let bindings in the current scope. */ + std::set let_defined_vars_; + /*! * \brief Update read/write buffers and regions with provided buffer and * region @@ -106,6 +124,28 @@ class MemoryAccessDetector : public StmtExprVisitor { regions->push_back(region); } + /*! + * \brief Update the set of read variables with the given variable + * \param var The variable to add to the set of read variables + */ + void UpdateReadVar(const Var &var) { + for (const auto &v : read_vars_) { + if (v.same_as(var)) return; + } + read_vars_.push_back(var); + } + + /*! + * \brief Update the set of write variables with the given variable + * \param var The variable to add to the set of write variables + */ + void UpdateWriteVar(const Var &var) { + for (const auto &v : write_vars_) { + if (v.same_as(var)) return; + } + write_vars_.push_back(var); + } + /*! * \brief Process a buffer region argument from reduce operation * \param arg The argument which could be BufferRegion, BufferLoad, or @@ -227,10 +267,19 @@ class MemoryAccessDetector : public StmtExprVisitor { void VisitStmt_(const LetStmtNode *op) override { let_bindings_[op->var.get()] = op->value; - StmtExprVisitor::VisitStmt_(op); + let_defined_vars_.insert(op->var.get()); + UpdateWriteVar(op->var); + StmtExprVisitor::VisitStmt(op->body); let_bindings_.erase(op->var.get()); } + void VisitExpr_(const VarNode *op) override { + if (let_defined_vars_.count(op)) { + UpdateReadVar(tvm::ffi::GetRef(op)); + } + StmtExprVisitor::VisitExpr_(op); + } + void VisitExpr_(const BufferLoadNode *op) override { std::vector relaxed_region; size_t num_indices = op->indices.size(); diff --git a/src/transform/if_condition_extract.cc b/src/transform/if_condition_extract.cc new file mode 100644 index 0000000000..9df3de18f5 --- /dev/null +++ b/src/transform/if_condition_extract.cc @@ -0,0 +1,123 @@ +/*! + * \file if_condition_extract.cc + * \brief Extract if conditions into temporary LetStmt variables, then expand if statements to all branches. + */ + +#include +#include +#include +#include +#include +#include + +#include "../op/builtin.h" + +namespace tvm { +namespace tl { + +using namespace tir; + +class IfConditionExtractor : public StmtExprMutator { +public: + static PrimFunc Substitute(PrimFunc &f) { + auto rewriter = IfConditionExtractor(); + f.CopyOnWrite()->body = rewriter(f->body); + return f; + } + +private: + IfConditionExtractor() = default; + + // counter to generate unique name for each IfStmt + int counter_ = 0; + + //! \brief Check if the expression is a simple variable. + bool IsSimpleVar(const PrimExpr &expr) { + return expr.as() != nullptr; + } + + Stmt VisitStmt_(const IfThenElseNode *op) final { + PrimExpr condition = VisitExpr(op->condition); + Stmt then_case = VisitStmt(op->then_case); + Optional else_case = op->else_case; + if (else_case.defined()) { + else_case = VisitStmt(else_case.value()); + } + + std::string var_name = "__cond_" + std::to_string(counter_++); + Var cond_var(var_name, DataType::Bool()); + bool is_simple = true; + + if (IsSimpleVar(condition)) { + // If the condition is already a simple variable, no need to extract it. + cond_var = Downcast(condition); + } else { + is_simple = false; + } + + auto bind_cond_var = [](const Stmt &sentence, const Var &cond) -> Stmt { + if (auto if_sentence = sentence.as()) { + PrimExpr new_cond = cond & if_sentence->condition; + return IfThenElse(new_cond, if_sentence->then_case, if_sentence->else_case); + } else { + return IfThenElse(cond, sentence); + } + }; + + auto bind_cond_var_body = [&](const Optional &body, const Var &cond) -> Stmt { + if (!body.defined()) { + return Stmt(); + } + if (auto seq = body.as()) { + Array new_seq; + for (auto sentence : seq->seq) { + new_seq.push_back(bind_cond_var(sentence, cond)); + } + return SeqStmt(std::move(new_seq)); + } else { + return bind_cond_var(body.value(), cond); + } + }; + + Array new_seq; + new_seq.insert(new_seq.end(), bind_cond_var_body(then_case, cond_var)); + if (else_case.defined()) new_seq.insert(new_seq.end(), bind_cond_var_body(else_case, cond_var)); + + Stmt body = new_seq.empty() ? Stmt() : (new_seq.size() == 1 ? new_seq[0] : SeqStmt(std::move(new_seq))); + if (is_simple) { + return body; + } else { + return LetStmt(cond_var, condition, body); + } + } + + Stmt VisitStmt_(const SeqStmtNode *op) final { + Array seq; + for (auto stmt : op->seq) { + auto new_stmt = VisitStmt(stmt); + if (!new_stmt.defined()) continue; + if (auto seq_node = new_stmt.as()) { + seq.insert(seq.end(), seq_node->seq.begin(), seq_node->seq.end()); + } else { + seq.push_back(new_stmt); + } + } + return SeqStmt(std::move(seq)); + } +}; + +using namespace tir::transform; +tvm::transform::Pass IfConditionExtract() { + auto pass_func = [=](PrimFunc f, const IRModule &m, const PassContext &ctx) { + return IfConditionExtractor::Substitute(f); + }; + return CreatePrimFuncPass(pass_func, 0, "tl.IfConditionExtract", {}); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef().def("tl.transform.IfConditionExtract", IfConditionExtract); +} + +} // namespace tl +} // namespace tvm diff --git a/tilelang/engine/phase.py b/tilelang/engine/phase.py index 8028d8caaa..b1b3e65624 100644 --- a/tilelang/engine/phase.py +++ b/tilelang/engine/phase.py @@ -196,9 +196,12 @@ def LowerAndLegalize(mod: IRModule, target: Target) -> IRModule: mod = tilelang.transform.Simplify()(mod) if allow_autoschedule(): # Auto schedule for high-level operations - mod = tilelang.transform.IfStmtBinding()(mod) + mod = tilelang.transform.IfConditionExtract()(mod) + # print("IfConditionExtract done") + # print(mod) mod = tilelang.transform.AutoSchedule(False)(mod) mod = tilelang.transform.Simplify()(mod) + print(mod) # Set layouts for reducers mod = tilelang.transform.LayoutReducer()(mod) # Infer memory layouts for fragments and shared memory @@ -330,4 +333,6 @@ def OptimizeForTarget(mod: IRModule, target: Target) -> IRModule: # Transform threadblock to persistent threadblock mod = tilelang.transform.PersistThreadblock()(mod) + print(mod) + return mod diff --git a/tilelang/transform/__init__.py b/tilelang/transform/__init__.py index 9e40425953..8a47b30543 100644 --- a/tilelang/transform/__init__.py +++ b/tilelang/transform/__init__.py @@ -544,6 +544,9 @@ def LowerSharedTmem(): """LowerSharedTmem""" return _ffi_api.LowerSharedTmem() # type: ignore +def IfConditionExtract(): + """Extract if condition to LetStmt variables.""" + return _ffi_api.IfConditionExtract() # type: ignore def AutoSchedule(enable_epi: bool): """Auto schedule for high-level operations""" From 2b1e779634d12bb0c798dbfdf38629cd63f89f3b Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Fri, 27 Mar 2026 17:34:14 +0800 Subject: [PATCH 2/6] Fix pipeline scheduling with Let defined variables. --- examples/example_flex_attn_wgmma.py | 647 ------------------ examples/example_flex_attn_wgmma_unchanged.py | 635 ----------------- src/transform/auto_schedule.cc | 513 +++++++++++--- src/transform/auto_schedule/ir_structure.h | 26 +- src/transform/auto_schedule/memory_detector.h | 9 +- tilelang/engine/phase.py | 2 +- 6 files changed, 444 insertions(+), 1388 deletions(-) delete mode 100644 examples/example_flex_attn_wgmma.py delete mode 100644 examples/example_flex_attn_wgmma_unchanged.py diff --git a/examples/example_flex_attn_wgmma.py b/examples/example_flex_attn_wgmma.py deleted file mode 100644 index 12accb7dd9..0000000000 --- a/examples/example_flex_attn_wgmma.py +++ /dev/null @@ -1,647 +0,0 @@ -import argparse -import itertools -from typing import Callable - -import tilelang -import tilelang.language as T -import torch - - -@tilelang.jit( - pass_configs={ - tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, - # tilelang.PassConfigKey.TL_ENABLE_AUTO_SCHEDULE: True - } -) -def flashattn_fwd( - batch: int, - heads: int, - dim_qk: int, - dim_vo: int, - softmax_scale: float, - mask_fn: Callable[[int, int, int, int, int, int, T.ptr], bool], - block_mask_fn: Callable[[int, int, int, int, int, int, int, int, T.ptr], bool], - block_qo: int = 64, - block_kv: int = 64, - num_stages: int = 2, - thread_num: int = 128, -): - scale = softmax_scale * 1.44269504 # log2(e) - dtype = T.bfloat16 - accum_dtype = T.float32 - - seq_len_qo = T.dynamic('seq_len_qo') - seq_len_kv = T.dynamic('seq_len_kv') - - @T.prim_func - def flash_fwd( - q_global: T.Tensor([batch, seq_len_qo, heads, dim_qk], dtype), - k_global: T.Tensor([batch, seq_len_kv, heads, dim_qk], dtype), - v_global: T.Tensor([batch, seq_len_kv, heads, dim_vo], dtype), - o_global: T.Tensor([batch, seq_len_qo, heads, dim_vo], dtype), - lse_global: T.Tensor([batch, heads, seq_len_qo], accum_dtype), - mask_custom_data: T.ptr, - block_mask_custom_data: T.ptr, - ): - with T.Kernel(T.ceildiv(seq_len_qo, block_qo), heads, batch, threads=thread_num) as ( - bx, - by, - bz, - ): - q_shared = T.alloc_shared([block_qo, dim_qk], dtype) - k_shared = T.alloc_shared([block_kv, dim_qk], dtype) - v_shared = T.alloc_shared([block_kv, dim_vo], dtype) - acc_s = T.alloc_fragment([block_qo, block_kv], accum_dtype) - acc_s_cast = T.alloc_fragment([block_qo, block_kv], dtype) - acc_o = T.alloc_fragment([block_qo, dim_vo], accum_dtype) - scores_max = T.alloc_fragment([block_qo], accum_dtype) - scores_max_prev = T.alloc_fragment([block_qo], accum_dtype) - scores_scale = T.alloc_fragment([block_qo], accum_dtype) - scores_sum = T.alloc_fragment([block_qo], accum_dtype) - scaled_sum_exp = T.alloc_fragment([block_qo], accum_dtype) - - T.copy(q_global[bz, bx * block_qo : (bx + 1) * block_qo, by, :], q_shared) - T.fill(acc_o, 0) - T.fill(scaled_sum_exp, 0) - T.fill(scores_max, -2.**100) - - for k in T.Pipelined(T.ceildiv(seq_len_kv, block_kv), num_stages=num_stages): - cond = block_mask_fn( - bz, - by, - bx * block_qo, - T.min(seq_len_qo, (bx + 1) * block_qo), - k * block_kv, - T.min(seq_len_kv, (k + 1) * block_kv), - seq_len_qo, - seq_len_kv, - block_mask_custom_data, - ) - if cond: - T.copy(k_global[bz, k * block_kv : (k + 1) * block_kv, by, :], k_shared) - for i, j in T.Parallel(block_qo, block_kv): - acc_s[i, j] = T.if_then_else( - mask_fn( - bz, - by, - bx * block_qo + i, - k * block_kv + j, - seq_len_qo, - seq_len_kv, - mask_custom_data, - ) - and k * block_kv + j < seq_len_kv, - 0, - -2.**100, - ) - T.gemm( - q_shared, - k_shared, - acc_s, - transpose_B=True, - policy=T.GemmWarpPolicy.FullRow, - ) - T.copy(v_global[bz, k * block_kv : (k + 1) * block_kv, by, :], v_shared) - T.copy(scores_max, scores_max_prev) - T.reduce_max(acc_s, scores_max, dim=1, clear=False) - for i in T.Parallel(block_qo): - scores_max[i] = T.max(scores_max[i], scores_max_prev[i]) - for i in T.Parallel(block_qo): - scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale) - for i, j in T.Parallel(block_qo, dim_vo): - acc_o[i, j] *= scores_scale[i] - for i, j in T.Parallel(block_qo, block_kv): - acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale) - T.copy(acc_s, acc_s_cast) - T.gemm(acc_s_cast, v_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) - T.reduce_sum(acc_s, scores_sum, dim=1) - for i in T.Parallel(block_qo): - scaled_sum_exp[i] = scaled_sum_exp[i] * scores_scale[i] + scores_sum[i] - - for i, j in T.Parallel(block_qo, dim_vo): - acc_o[i, j] /= scaled_sum_exp[i] - T.copy(acc_o, o_global[bz, bx * block_qo : (bx + 1) * block_qo, by, :]) - - for i in T.Parallel(block_qo): - scaled_sum_exp[i] = T.log2(scaled_sum_exp[i]) + scores_max[i] * scale - T.copy(scaled_sum_exp, lse_global[bz, by, bx * block_qo : (bx + 1) * block_qo]) - - return flash_fwd - - -@tilelang.jit( - pass_configs={ - tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, - # tilelang.PassConfigKey.TL_ENABLE_AUTO_SCHEDULE: True - }, -) -def flashattn_bwd_preprocess(batch, heads, dim): - dtype = T.bfloat16 - accum_dtype = T.float32 - seq_len = T.dynamic('seq_len') - shape = [batch, seq_len, heads, dim] - blk = 32 - - @T.prim_func - def flash_bwd_prep( - o_global: T.Tensor(shape, dtype), - grad_o_global: T.Tensor(shape, dtype), - delta_global: T.Tensor([batch, heads, seq_len], accum_dtype), - ): - with T.Kernel(heads, T.ceildiv(seq_len, blk), batch) as (bx, by, bz): - o = T.alloc_fragment([blk, blk], dtype) - do = T.alloc_fragment([blk, blk], dtype) - acc = T.alloc_fragment([blk, blk], accum_dtype) - delta = T.alloc_fragment([blk], accum_dtype) - T.clear(acc) - for k in range(T.ceildiv(dim, blk)): - T.copy( - o_global[ - bz, - by * blk : (by + 1) * blk, - bx, - k * blk : (k + 1) * blk, - ], - o, - ) - T.copy( - grad_o_global[ - bz, - by * blk : (by + 1) * blk, - bx, - k * blk : (k + 1) * blk, - ], - do, - ) - for i, j in T.Parallel(blk, blk): - acc[i, j] += o[i, j] * do[i, j] - T.reduce_sum(acc, delta, 1) - T.copy(delta, delta_global[bz, bx, by * blk : (by + 1) * blk]) - - return flash_bwd_prep - - -@tilelang.jit( - pass_configs={ - tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, - # tilelang.PassConfigKey.TL_ENABLE_AUTO_SCHEDULE: True - } -) -def flashattn_bwd( - batch: int, - heads: int, - dim_qk: int, - dim_vo: int, - softmax_scale: float, - mask_fn: Callable[[int, int, int, int, int, int], bool], - block_mask_fn: Callable[[int, int, int, int, int, int, int, int], bool], - block_qo: int = 32, - block_kv: int = 64, - num_stages: int = 2, - thread_num: int = 128, -): - scale = softmax_scale * 1.44269504 # log2(e) - dtype = T.bfloat16 - accum_dtype = T.float32 - - seq_len_qo = T.dynamic('seq_len_qo') - seq_len_kv = T.dynamic('seq_len_kv') - - @T.prim_func - def flash_bwd( - q_global: T.Tensor([batch, seq_len_qo, heads, dim_qk], dtype), - k_global: T.Tensor([batch, seq_len_kv, heads, dim_qk], dtype), - v_global: T.Tensor([batch, seq_len_kv, heads, dim_vo], dtype), - grad_o_global: T.Tensor([batch, seq_len_qo, heads, dim_vo], dtype), - lse_global: T.Tensor([batch, heads, seq_len_qo], accum_dtype), - delta_global: T.Tensor([batch, heads, seq_len_qo], accum_dtype), - grad_q_global: T.Tensor( - [batch, heads, T.ceildiv(seq_len_qo, block_qo) * block_qo, dim_qk], accum_dtype - ), - grad_k_global: T.Tensor([batch, seq_len_kv, heads, dim_qk], dtype), - grad_v_global: T.Tensor([batch, seq_len_kv, heads, dim_vo], dtype), - mask_custom_data: T.ptr, - block_mask_custom_data: T.ptr, - ): - with T.Kernel(heads, T.ceildiv(seq_len_kv, block_kv), batch, threads=thread_num) as ( - bx, - by, - bz, - ): - T.disable_warp_group_reg_alloc() - K_shared = T.alloc_shared([block_kv, dim_qk], dtype) - dsT_shared = T.alloc_shared([block_kv, block_qo], dtype) - q = T.alloc_shared([block_qo, dim_qk], dtype) - V_shared = T.alloc_shared([block_kv, dim_vo], dtype) - qkT = T.alloc_fragment([block_kv, block_qo], accum_dtype) - dsT = T.alloc_fragment([block_kv, block_qo], accum_dtype) - qkT_cast = T.alloc_shared([block_kv, block_qo], dtype) - lse_shared = T.alloc_shared([block_qo], accum_dtype) - delta = T.alloc_shared([block_qo], accum_dtype) - do = T.alloc_shared([block_qo, dim_vo], dtype) - dv = T.alloc_fragment([block_kv, dim_vo], accum_dtype) - dk = T.alloc_fragment([block_kv, dim_qk], accum_dtype) - dq = T.alloc_fragment([block_qo, dim_qk], accum_dtype) - dv_shared = T.alloc_shared([block_kv, dim_vo], dtype) - dk_shared = T.alloc_shared([block_kv, dim_qk], dtype) - dq_shared = T.alloc_shared([block_qo, dim_qk], accum_dtype) - - T.copy(k_global[bz, by * block_kv : (by + 1) * block_kv, bx, :], K_shared) - T.copy(v_global[bz, by * block_kv : (by + 1) * block_kv, bx, :], V_shared) - T.clear(dv) - T.clear(dk) - for k in T.Pipelined(0, T.ceildiv(seq_len_qo, block_qo), num_stages=num_stages): - if block_mask_fn( - bz, - bx, - k * block_qo, - T.min(seq_len_qo, (k + 1) * block_qo), - by * block_kv, - T.min(seq_len_kv, (by + 1) * block_kv), - seq_len_qo, - seq_len_kv, - block_mask_custom_data, - ): - T.copy(q_global[bz, k * block_qo : (k + 1) * block_qo, bx, :], q) - T.clear(qkT) - T.gemm(K_shared, q, qkT, transpose_B=True) - T.copy(grad_o_global[bz, k * block_qo : (k + 1) * block_qo, bx, :], do) - T.clear(dsT) - T.gemm(V_shared, do, dsT, transpose_B=True) - - T.copy(lse_global[bz, bx, k * block_qo : (k + 1) * block_qo], lse_shared) - for i, j in T.Parallel(block_kv, block_qo): - qkT[i, j] = T.exp2(qkT[i, j] * scale - lse_shared[j]) - # We don't need to handle OOB positions, - # since OOB values won't affect other positions here. - for i, j in T.Parallel(block_kv, block_qo): - qkT[i, j] = T.if_then_else( - mask_fn( - bz, - bx, - k * block_qo + j, - by * block_kv + i, - seq_len_qo, - seq_len_kv, - mask_custom_data, - ), - qkT[i, j], - 0, - ) - T.copy(qkT, qkT_cast) - T.gemm(qkT_cast, do, dv) - - T.copy(delta_global[bz, bx, k * block_qo : (k + 1) * block_qo], delta) - - for i, j in T.Parallel(block_kv, block_qo): - dsT_shared[i, j] = qkT[i, j] * (dsT[i, j] - delta[j]) * softmax_scale - T.gemm(dsT_shared, q, dk) - - T.clear(dq) - T.gemm(dsT_shared, K_shared, dq, transpose_A=True) - T.copy(dq, dq_shared) - T.atomic_add( - grad_q_global[bz, bx, k * block_qo : (k + 1) * block_qo, :], - dq_shared, - ) - T.copy(dv, dv_shared) - T.copy(dk, dk_shared) - T.copy(dv_shared, grad_v_global[bz, by * block_kv : (by + 1) * block_kv, bx, :]) - T.copy(dk_shared, grad_k_global[bz, by * block_kv : (by + 1) * block_kv, bx, :]) - - return flash_bwd - - -class _FlexAttn(torch.autograd.Function): - @staticmethod - def forward( - ctx: torch.autograd.Function, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - softmax_scale: float, - mask_fn: Callable[[int, int], bool], - block_mask_fn: Callable[[int, int, int, int], bool], - mask_custom_data: torch.Tensor, - block_mask_custom_data: torch.Tensor, - ): - batch, seq_len_qo, head, dim_qk = q.shape - _, seq_len_kv, _, dim_vo = v.shape - o = torch.empty((batch, seq_len_qo, head, dim_vo), dtype=q.dtype, device=q.device) - lse = torch.empty((batch, head, seq_len_qo), dtype=torch.float32, device=q.device) - mod = flashattn_fwd( - batch, - head, - dim_qk, - dim_vo, - softmax_scale, - mask_fn, - block_mask_fn, - ) - mod( - q, - k, - v, - o, - lse, - mask_custom_data, - block_mask_custom_data, - ) - ctx.save_for_backward( - q, - k, - v, - o, - lse, - mask_custom_data, - block_mask_custom_data, - ) - ctx.softmax_scale = softmax_scale - ctx.mask_fn = mask_fn - ctx.block_mask_fn = block_mask_fn - return o, lse - - @staticmethod - def backward(ctx: torch.autograd.Function, do: torch.Tensor, dlse: None): - ( - q, - k, - v, - o, - lse, - mask_custom_data, - block_mask_custom_data, - ) = ctx.saved_tensors - batch, seq_len_qo, head, dim_qk = q.shape - _, seq_len_kv, _, dim_vo = v.shape - - do, q, k, v, o = [x.contiguous() for x in (do, q, k, v, o)] - - delta = torch.empty((batch, head, seq_len_qo), dtype=torch.float32, device=q.device) - flashattn_bwd_preprocess(batch, head, dim_vo)(o, do, delta) - - block_qo = 32 - mod = flashattn_bwd( - batch, - head, - dim_qk, - dim_vo, - ctx.softmax_scale, - ctx.mask_fn, - ctx.block_mask_fn, - block_qo=block_qo, - ) - dq = torch.zeros( - (batch, head, (seq_len_qo + block_qo - 1) // block_qo * block_qo, dim_qk), - device=q.device, - dtype=torch.float32, - ) - dk = torch.zeros_like(k) - dv = torch.empty_like(v) - mod( - q, - k, - v, - do, - lse, - delta, - dq, - dk, - dv, - mask_custom_data, - block_mask_custom_data, - ) - dq = dq[:, :, :seq_len_qo, :].transpose(1, 2).bfloat16() # TODO: tilelang cast - return dq, dk, dv, *([None] * 5) - - -flex_attn = _FlexAttn.apply - - -def ref_program( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - softmax_scale: float, - bool_mask: torch.Tensor, -): - score = torch.einsum('bshd,bthd->bhst', q, k) * softmax_scale - score = score.masked_fill(~bool_mask, float('-inf')) - p = torch.softmax(score, dim=-1) - lse = torch.logsumexp(score, dim=-1) * 1.44269504 - o = torch.einsum('bhst,bthd->bshd', p, v) - return o, lse - - -def main( - BATCH: int, - H: int, - N_CTX_QO: int, - D_HEAD_QK: int, - N_CTX_KV: int, - D_HEAD_VO: int, - mode: str, -): - if mode == 'causal': - bool_mask = ( - (N_CTX_QO - torch.arange(N_CTX_QO, device='cuda'))[:, None] - <= (N_CTX_KV - torch.arange(N_CTX_KV, device='cuda'))[None, :] - ).expand(BATCH, H, N_CTX_QO, N_CTX_KV) - - @T.macro - def mask_fn( - b, - h, - q_idx, - k_idx, - seq_len_qo, - seq_len_kv, - mask_custom_data, - ): - return (seq_len_qo - q_idx) <= (seq_len_kv - k_idx) - - @T.macro - def block_mask_fn( - b, - h, - q_start, - q_end, - k_start, - k_end, - seq_len_qo, - seq_len_kv, - block_mask_custom_data, - ): - return (seq_len_qo - q_end - 1) <= (seq_len_kv - k_start) - - mask_custom_data = None - block_mask_custom_data = None - - elif mode == 'full': - bool_mask = torch.ones((BATCH, H, N_CTX_QO, N_CTX_KV), dtype=torch.bool, device='cuda') - - @T.macro - def mask_fn( - b, - h, - q_idx, - k_idx, - seq_len_qo, - seq_len_kv, - mask_custom_data, - ): - return True - - @T.macro - def block_mask_fn( - b, - h, - q_start, - q_end, - k_start, - k_end, - seq_len_qo, - seq_len_kv, - block_mask_custom_data, - ): - return True - - mask_custom_data = None - block_mask_custom_data = None - elif mode == 'random': - bool_mask = torch.rand((BATCH, H, N_CTX_QO, N_CTX_KV), device='cuda') < 0.5 - - @T.macro - def mask_fn( - b, - h, - q_idx, - k_idx, - seq_len_qo, - seq_len_kv, - mask_custom_data, - ): - mask = T.make_tensor( - mask_custom_data, - shape=(BATCH, H, seq_len_qo, seq_len_kv), - dtype=T.bool, - ) - return mask[b, h, q_idx, k_idx] - - @T.macro - def block_mask_fn( - b, - h, - q_start, - q_end, - k_start, - k_end, - seq_len_qo, - seq_len_kv, - block_mask_custom_data, - ): - block_mask = T.make_tensor( - block_mask_custom_data, - shape=(BATCH, H, T.ceildiv(seq_len_qo, 128), T.ceildiv(seq_len_kv, 128)), - dtype=T.bool, - ) - return T.if_then_else( - q_start // 128 == (q_end - 1) // 128 and k_start // 128 == (k_end - 1) // 128, - block_mask[b, h, q_start // 128, k_start // 128], - True, - ) - - mask_custom_data = bool_mask - block_mask_custom_data = ( - bool_mask.unflatten(-2, (-1, 128)).unflatten(-1, (-1, 128)).any(dim=(-1, -3)) - ) - - Q = ( - torch.empty(BATCH, N_CTX_QO, H, D_HEAD_QK, dtype=torch.bfloat16, device='cuda') - .normal_() - .requires_grad_() - ) - K = ( - torch.empty(BATCH, N_CTX_KV, H, D_HEAD_QK, dtype=torch.bfloat16, device='cuda') - .normal_() - .requires_grad_() - ) - V = ( - torch.empty(BATCH, N_CTX_KV, H, D_HEAD_VO, dtype=torch.bfloat16, device='cuda') - .normal_() - .requires_grad_() - ) - dO = ( - torch.empty(BATCH, N_CTX_QO, H, D_HEAD_VO, dtype=torch.bfloat16, device='cuda') - .normal_() - .requires_grad_() - ) - O, lse = flex_attn( - Q, - K, - V, - D_HEAD_QK**-0.5, - mask_fn, - block_mask_fn, - mask_custom_data, - block_mask_custom_data, - ) - O.backward(dO) - dQ, Q.grad = Q.grad.clone(), None - dK, K.grad = K.grad.clone(), None - dV, V.grad = V.grad.clone(), None - - Q_ref = Q.float().detach().requires_grad_() - K_ref = K.float().detach().requires_grad_() - V_ref = V.float().detach().requires_grad_() - O_ref, lse_ref = ref_program(Q_ref, K_ref, V_ref, D_HEAD_QK**-0.5, bool_mask) - O_ref.backward(dO) - dQ_ref, Q_ref.grad = Q_ref.grad.clone(), None - dK_ref, K_ref.grad = K_ref.grad.clone(), None - dV_ref, V_ref.grad = V_ref.grad.clone(), None - - torch.testing.assert_close(lse, lse_ref, rtol=3e-6, atol=1e-5) - torch.testing.assert_close(O.float(), O_ref, rtol=1e-2, atol=1e-2) - torch.testing.assert_close(dV.float(), dV_ref, rtol=1e-2, atol=1e-2) - torch.testing.assert_close(dK.float(), dK_ref, rtol=1e-2, atol=1e-2) - torch.testing.assert_close(dQ.float(), dQ_ref, rtol=1e-2, atol=1e-2) - print('All checks passed.✅') - - def run(): - Q.grad = K.grad = V.grad = None - flex_attn( - Q, - K, - V, - D_HEAD_QK**-0.5, - mask_fn, - block_mask_fn, - mask_custom_data, - block_mask_custom_data, - )[0].backward(dO) - - with torch.profiler.profile() as prof: - for _ in range(5): - run() - print(prof.key_averages().table(sort_by='cuda_time_total', row_limit=10)) - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--batch', type=int, default=4, help='Batch size') - parser.add_argument('--h', type=int, default=16, help='Number of heads') - parser.add_argument('--n_ctx_qo', type=int, default=4096, help='Context size') - parser.add_argument('--d_head_qk', type=int, default=192, help='Head dimension') - parser.add_argument('--n_ctx_kv', type=int, default=8192, help='Context size') - parser.add_argument('--d_head_vo', type=int, default=128, help='Head dimension') - parser.add_argument('--mask_mode', type=str, default='random', help='Causal flag') - args = parser.parse_args() - main( - args.batch, - args.h, - args.n_ctx_qo, - args.d_head_qk, - args.n_ctx_kv, - args.d_head_vo, - args.mask_mode, - ) diff --git a/examples/example_flex_attn_wgmma_unchanged.py b/examples/example_flex_attn_wgmma_unchanged.py deleted file mode 100644 index fe4fdae8ad..0000000000 --- a/examples/example_flex_attn_wgmma_unchanged.py +++ /dev/null @@ -1,635 +0,0 @@ -import argparse -import itertools -from typing import Callable - -import tilelang -import tilelang.language as T -import torch - - -@tilelang.jit(pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True}) -def flashattn_fwd( - batch: int, - heads: int, - dim_qk: int, - dim_vo: int, - softmax_scale: float, - mask_fn: Callable[[int, int, int, int, int, int, T.ptr], bool], - block_mask_fn: Callable[[int, int, int, int, int, int, int, int, T.ptr], bool], - block_qo: int = 64, - block_kv: int = 64, - num_stages: int = 2, - thread_num: int = 128, -): - scale = softmax_scale * 1.44269504 # log2(e) - dtype = T.bfloat16 - accum_dtype = T.float32 - - seq_len_qo = T.dynamic('seq_len_qo') - seq_len_kv = T.dynamic('seq_len_kv') - - @T.prim_func - def flash_fwd( - q_global: T.Tensor([batch, seq_len_qo, heads, dim_qk], dtype), - k_global: T.Tensor([batch, seq_len_kv, heads, dim_qk], dtype), - v_global: T.Tensor([batch, seq_len_kv, heads, dim_vo], dtype), - o_global: T.Tensor([batch, seq_len_qo, heads, dim_vo], dtype), - lse_global: T.Tensor([batch, heads, seq_len_qo], accum_dtype), - mask_custom_data: T.ptr, - block_mask_custom_data: T.ptr, - ): - with T.Kernel(T.ceildiv(seq_len_qo, block_qo), heads, batch, threads=thread_num) as ( - bx, - by, - bz, - ): - q_shared = T.alloc_shared([block_qo, dim_qk], dtype) - k_shared = T.alloc_shared([block_kv, dim_qk], dtype) - v_shared = T.alloc_shared([block_kv, dim_vo], dtype) - acc_s = T.alloc_fragment([block_qo, block_kv], accum_dtype) - acc_s_cast = T.alloc_fragment([block_qo, block_kv], dtype) - acc_o = T.alloc_fragment([block_qo, dim_vo], accum_dtype) - scores_max = T.alloc_fragment([block_qo], accum_dtype) - scores_max_prev = T.alloc_fragment([block_qo], accum_dtype) - scores_scale = T.alloc_fragment([block_qo], accum_dtype) - scores_sum = T.alloc_fragment([block_qo], accum_dtype) - scaled_sum_exp = T.alloc_fragment([block_qo], accum_dtype) - - T.copy(q_global[bz, bx * block_qo : (bx + 1) * block_qo, by, :], q_shared) - T.fill(acc_o, 0) - T.fill(scaled_sum_exp, 0) - T.fill(scores_max, -2.**100) - - for k in T.Pipelined(T.ceildiv(seq_len_kv, block_kv), num_stages=num_stages): - if block_mask_fn( - bz, - by, - bx * block_qo, - T.min(seq_len_qo, (bx + 1) * block_qo), - k * block_kv, - T.min(seq_len_kv, (k + 1) * block_kv), - seq_len_qo, - seq_len_kv, - block_mask_custom_data, - ): - T.copy(k_global[bz, k * block_kv : (k + 1) * block_kv, by, :], k_shared) - for i, j in T.Parallel(block_qo, block_kv): - acc_s[i, j] = T.if_then_else( - mask_fn( - bz, - by, - bx * block_qo + i, - k * block_kv + j, - seq_len_qo, - seq_len_kv, - mask_custom_data, - ) - and k * block_kv + j < seq_len_kv, - 0, - -2.**100, - ) - T.gemm( - q_shared, - k_shared, - acc_s, - transpose_B=True, - policy=T.GemmWarpPolicy.FullRow, - ) - T.copy(v_global[bz, k * block_kv : (k + 1) * block_kv, by, :], v_shared) - T.copy(scores_max, scores_max_prev) - T.reduce_max(acc_s, scores_max, dim=1, clear=False) - for i in T.Parallel(block_qo): - scores_max[i] = T.max(scores_max[i], scores_max_prev[i]) - for i in T.Parallel(block_qo): - scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale) - for i, j in T.Parallel(block_qo, dim_vo): - acc_o[i, j] *= scores_scale[i] - for i, j in T.Parallel(block_qo, block_kv): - acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale) - T.copy(acc_s, acc_s_cast) - T.gemm(acc_s_cast, v_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) - T.reduce_sum(acc_s, scores_sum, dim=1) - for i in T.Parallel(block_qo): - scaled_sum_exp[i] = scaled_sum_exp[i] * scores_scale[i] + scores_sum[i] - - for i, j in T.Parallel(block_qo, dim_vo): - acc_o[i, j] /= scaled_sum_exp[i] - T.copy(acc_o, o_global[bz, bx * block_qo : (bx + 1) * block_qo, by, :]) - - for i in T.Parallel(block_qo): - scaled_sum_exp[i] = T.log2(scaled_sum_exp[i]) + scores_max[i] * scale - T.copy(scaled_sum_exp, lse_global[bz, by, bx * block_qo : (bx + 1) * block_qo]) - - return flash_fwd - - -@tilelang.jit( - pass_configs={ - tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, - }, -) -def flashattn_bwd_preprocess(batch, heads, dim): - dtype = T.bfloat16 - accum_dtype = T.float32 - seq_len = T.dynamic('seq_len') - shape = [batch, seq_len, heads, dim] - blk = 32 - - @T.prim_func - def flash_bwd_prep( - o_global: T.Tensor(shape, dtype), - grad_o_global: T.Tensor(shape, dtype), - delta_global: T.Tensor([batch, heads, seq_len], accum_dtype), - ): - with T.Kernel(heads, T.ceildiv(seq_len, blk), batch) as (bx, by, bz): - o = T.alloc_fragment([blk, blk], dtype) - do = T.alloc_fragment([blk, blk], dtype) - acc = T.alloc_fragment([blk, blk], accum_dtype) - delta = T.alloc_fragment([blk], accum_dtype) - T.clear(acc) - for k in range(T.ceildiv(dim, blk)): - T.copy( - o_global[ - bz, - by * blk : (by + 1) * blk, - bx, - k * blk : (k + 1) * blk, - ], - o, - ) - T.copy( - grad_o_global[ - bz, - by * blk : (by + 1) * blk, - bx, - k * blk : (k + 1) * blk, - ], - do, - ) - for i, j in T.Parallel(blk, blk): - acc[i, j] += o[i, j] * do[i, j] - T.reduce_sum(acc, delta, 1) - T.copy(delta, delta_global[bz, bx, by * blk : (by + 1) * blk]) - - return flash_bwd_prep - - -@tilelang.jit(pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True}) -def flashattn_bwd( - batch: int, - heads: int, - dim_qk: int, - dim_vo: int, - softmax_scale: float, - mask_fn: Callable[[int, int, int, int, int, int], bool], - block_mask_fn: Callable[[int, int, int, int, int, int, int, int], bool], - block_qo: int = 32, - block_kv: int = 64, - num_stages: int = 2, - thread_num: int = 128, -): - scale = softmax_scale * 1.44269504 # log2(e) - dtype = T.bfloat16 - accum_dtype = T.float32 - - seq_len_qo = T.dynamic('seq_len_qo') - seq_len_kv = T.dynamic('seq_len_kv') - - @T.prim_func - def flash_bwd( - q_global: T.Tensor([batch, seq_len_qo, heads, dim_qk], dtype), - k_global: T.Tensor([batch, seq_len_kv, heads, dim_qk], dtype), - v_global: T.Tensor([batch, seq_len_kv, heads, dim_vo], dtype), - grad_o_global: T.Tensor([batch, seq_len_qo, heads, dim_vo], dtype), - lse_global: T.Tensor([batch, heads, seq_len_qo], accum_dtype), - delta_global: T.Tensor([batch, heads, seq_len_qo], accum_dtype), - grad_q_global: T.Tensor( - [batch, heads, T.ceildiv(seq_len_qo, block_qo) * block_qo, dim_qk], accum_dtype - ), - grad_k_global: T.Tensor([batch, seq_len_kv, heads, dim_qk], dtype), - grad_v_global: T.Tensor([batch, seq_len_kv, heads, dim_vo], dtype), - mask_custom_data: T.ptr, - block_mask_custom_data: T.ptr, - ): - with T.Kernel(heads, T.ceildiv(seq_len_kv, block_kv), batch, threads=thread_num) as ( - bx, - by, - bz, - ): - T.disable_warp_group_reg_alloc() - K_shared = T.alloc_shared([block_kv, dim_qk], dtype) - dsT_shared = T.alloc_shared([block_kv, block_qo], dtype) - q = T.alloc_shared([block_qo, dim_qk], dtype) - V_shared = T.alloc_shared([block_kv, dim_vo], dtype) - qkT = T.alloc_fragment([block_kv, block_qo], accum_dtype) - dsT = T.alloc_fragment([block_kv, block_qo], accum_dtype) - qkT_cast = T.alloc_shared([block_kv, block_qo], dtype) - lse_shared = T.alloc_shared([block_qo], accum_dtype) - delta = T.alloc_shared([block_qo], accum_dtype) - do = T.alloc_shared([block_qo, dim_vo], dtype) - dv = T.alloc_fragment([block_kv, dim_vo], accum_dtype) - dk = T.alloc_fragment([block_kv, dim_qk], accum_dtype) - dq = T.alloc_fragment([block_qo, dim_qk], accum_dtype) - dv_shared = T.alloc_shared([block_kv, dim_vo], dtype) - dk_shared = T.alloc_shared([block_kv, dim_qk], dtype) - dq_shared = T.alloc_shared([block_qo, dim_qk], accum_dtype) - - T.copy(k_global[bz, by * block_kv : (by + 1) * block_kv, bx, :], K_shared) - T.copy(v_global[bz, by * block_kv : (by + 1) * block_kv, bx, :], V_shared) - T.clear(dv) - T.clear(dk) - for k in T.Pipelined(0, T.ceildiv(seq_len_qo, block_qo), num_stages=num_stages): - if block_mask_fn( - bz, - bx, - k * block_qo, - T.min(seq_len_qo, (k + 1) * block_qo), - by * block_kv, - T.min(seq_len_kv, (by + 1) * block_kv), - seq_len_qo, - seq_len_kv, - block_mask_custom_data, - ): - T.copy(q_global[bz, k * block_qo : (k + 1) * block_qo, bx, :], q) - T.clear(qkT) - T.gemm(K_shared, q, qkT, transpose_B=True) - T.copy(grad_o_global[bz, k * block_qo : (k + 1) * block_qo, bx, :], do) - T.clear(dsT) - T.gemm(V_shared, do, dsT, transpose_B=True) - - T.copy(lse_global[bz, bx, k * block_qo : (k + 1) * block_qo], lse_shared) - for i, j in T.Parallel(block_kv, block_qo): - qkT[i, j] = T.exp2(qkT[i, j] * scale - lse_shared[j]) - # We don't need to handle OOB positions, - # since OOB values won't affect other positions here. - for i, j in T.Parallel(block_kv, block_qo): - qkT[i, j] = T.if_then_else( - mask_fn( - bz, - bx, - k * block_qo + j, - by * block_kv + i, - seq_len_qo, - seq_len_kv, - mask_custom_data, - ), - qkT[i, j], - 0, - ) - T.copy(qkT, qkT_cast) - T.gemm(qkT_cast, do, dv) - - T.copy(delta_global[bz, bx, k * block_qo : (k + 1) * block_qo], delta) - - for i, j in T.Parallel(block_kv, block_qo): - dsT_shared[i, j] = qkT[i, j] * (dsT[i, j] - delta[j]) * softmax_scale - T.gemm(dsT_shared, q, dk) - - T.clear(dq) - T.gemm(dsT_shared, K_shared, dq, transpose_A=True) - T.copy(dq, dq_shared) - T.atomic_add( - grad_q_global[bz, bx, k * block_qo : (k + 1) * block_qo, :], - dq_shared, - ) - T.copy(dv, dv_shared) - T.copy(dk, dk_shared) - T.copy(dv_shared, grad_v_global[bz, by * block_kv : (by + 1) * block_kv, bx, :]) - T.copy(dk_shared, grad_k_global[bz, by * block_kv : (by + 1) * block_kv, bx, :]) - - return flash_bwd - - -class _FlexAttn(torch.autograd.Function): - @staticmethod - def forward( - ctx: torch.autograd.Function, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - softmax_scale: float, - mask_fn: Callable[[int, int], bool], - block_mask_fn: Callable[[int, int, int, int], bool], - mask_custom_data: torch.Tensor, - block_mask_custom_data: torch.Tensor, - ): - batch, seq_len_qo, head, dim_qk = q.shape - _, seq_len_kv, _, dim_vo = v.shape - o = torch.empty((batch, seq_len_qo, head, dim_vo), dtype=q.dtype, device=q.device) - lse = torch.empty((batch, head, seq_len_qo), dtype=torch.float32, device=q.device) - mod = flashattn_fwd( - batch, - head, - dim_qk, - dim_vo, - softmax_scale, - mask_fn, - block_mask_fn, - ) - mod( - q, - k, - v, - o, - lse, - mask_custom_data, - block_mask_custom_data, - ) - ctx.save_for_backward( - q, - k, - v, - o, - lse, - mask_custom_data, - block_mask_custom_data, - ) - ctx.softmax_scale = softmax_scale - ctx.mask_fn = mask_fn - ctx.block_mask_fn = block_mask_fn - return o, lse - - @staticmethod - def backward(ctx: torch.autograd.Function, do: torch.Tensor, dlse: None): - ( - q, - k, - v, - o, - lse, - mask_custom_data, - block_mask_custom_data, - ) = ctx.saved_tensors - batch, seq_len_qo, head, dim_qk = q.shape - _, seq_len_kv, _, dim_vo = v.shape - - do, q, k, v, o = [x.contiguous() for x in (do, q, k, v, o)] - - delta = torch.empty((batch, head, seq_len_qo), dtype=torch.float32, device=q.device) - flashattn_bwd_preprocess(batch, head, dim_vo)(o, do, delta) - - block_qo = 32 - mod = flashattn_bwd( - batch, - head, - dim_qk, - dim_vo, - ctx.softmax_scale, - ctx.mask_fn, - ctx.block_mask_fn, - block_qo=block_qo, - ) - dq = torch.zeros( - (batch, head, (seq_len_qo + block_qo - 1) // block_qo * block_qo, dim_qk), - device=q.device, - dtype=torch.float32, - ) - dk = torch.zeros_like(k) - dv = torch.empty_like(v) - mod( - q, - k, - v, - do, - lse, - delta, - dq, - dk, - dv, - mask_custom_data, - block_mask_custom_data, - ) - dq = dq[:, :, :seq_len_qo, :].transpose(1, 2).bfloat16() # TODO: tilelang cast - return dq, dk, dv, *([None] * 5) - - -flex_attn = _FlexAttn.apply - - -def ref_program( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - softmax_scale: float, - bool_mask: torch.Tensor, -): - score = torch.einsum('bshd,bthd->bhst', q, k) * softmax_scale - score = score.masked_fill(~bool_mask, float('-inf')) - p = torch.softmax(score, dim=-1) - lse = torch.logsumexp(score, dim=-1) * 1.44269504 - o = torch.einsum('bhst,bthd->bshd', p, v) - return o, lse - - -def main( - BATCH: int, - H: int, - N_CTX_QO: int, - D_HEAD_QK: int, - N_CTX_KV: int, - D_HEAD_VO: int, - mode: str, -): - if mode == 'causal': - bool_mask = ( - (N_CTX_QO - torch.arange(N_CTX_QO, device='cuda'))[:, None] - <= (N_CTX_KV - torch.arange(N_CTX_KV, device='cuda'))[None, :] - ).expand(BATCH, H, N_CTX_QO, N_CTX_KV) - - @T.macro - def mask_fn( - b, - h, - q_idx, - k_idx, - seq_len_qo, - seq_len_kv, - mask_custom_data, - ): - return (seq_len_qo - q_idx) <= (seq_len_kv - k_idx) - - @T.macro - def block_mask_fn( - b, - h, - q_start, - q_end, - k_start, - k_end, - seq_len_qo, - seq_len_kv, - block_mask_custom_data, - ): - return (seq_len_qo - q_end - 1) <= (seq_len_kv - k_start) - - mask_custom_data = None - block_mask_custom_data = None - - elif mode == 'full': - bool_mask = torch.ones((BATCH, H, N_CTX_QO, N_CTX_KV), dtype=torch.bool, device='cuda') - - @T.macro - def mask_fn( - b, - h, - q_idx, - k_idx, - seq_len_qo, - seq_len_kv, - mask_custom_data, - ): - return True - - @T.macro - def block_mask_fn( - b, - h, - q_start, - q_end, - k_start, - k_end, - seq_len_qo, - seq_len_kv, - block_mask_custom_data, - ): - return True - - mask_custom_data = None - block_mask_custom_data = None - elif mode == 'random': - bool_mask = torch.rand((BATCH, H, N_CTX_QO, N_CTX_KV), device='cuda') < 0.5 - - @T.macro - def mask_fn( - b, - h, - q_idx, - k_idx, - seq_len_qo, - seq_len_kv, - mask_custom_data, - ): - mask = T.make_tensor( - mask_custom_data, - shape=(BATCH, H, seq_len_qo, seq_len_kv), - dtype=T.bool, - ) - return mask[b, h, q_idx, k_idx] - - @T.macro - def block_mask_fn( - b, - h, - q_start, - q_end, - k_start, - k_end, - seq_len_qo, - seq_len_kv, - block_mask_custom_data, - ): - block_mask = T.make_tensor( - block_mask_custom_data, - shape=(BATCH, H, T.ceildiv(seq_len_qo, 128), T.ceildiv(seq_len_kv, 128)), - dtype=T.bool, - ) - return T.if_then_else( - q_start // 128 == (q_end - 1) // 128 and k_start // 128 == (k_end - 1) // 128, - block_mask[b, h, q_start // 128, k_start // 128], - True, - ) - - mask_custom_data = bool_mask - block_mask_custom_data = ( - bool_mask.unflatten(-2, (-1, 128)).unflatten(-1, (-1, 128)).any(dim=(-1, -3)) - ) - - Q = ( - torch.empty(BATCH, N_CTX_QO, H, D_HEAD_QK, dtype=torch.bfloat16, device='cuda') - .normal_() - .requires_grad_() - ) - K = ( - torch.empty(BATCH, N_CTX_KV, H, D_HEAD_QK, dtype=torch.bfloat16, device='cuda') - .normal_() - .requires_grad_() - ) - V = ( - torch.empty(BATCH, N_CTX_KV, H, D_HEAD_VO, dtype=torch.bfloat16, device='cuda') - .normal_() - .requires_grad_() - ) - dO = ( - torch.empty(BATCH, N_CTX_QO, H, D_HEAD_VO, dtype=torch.bfloat16, device='cuda') - .normal_() - .requires_grad_() - ) - O, lse = flex_attn( - Q, - K, - V, - D_HEAD_QK**-0.5, - mask_fn, - block_mask_fn, - mask_custom_data, - block_mask_custom_data, - ) - O.backward(dO) - dQ, Q.grad = Q.grad.clone(), None - dK, K.grad = K.grad.clone(), None - dV, V.grad = V.grad.clone(), None - - Q_ref = Q.float().detach().requires_grad_() - K_ref = K.float().detach().requires_grad_() - V_ref = V.float().detach().requires_grad_() - O_ref, lse_ref = ref_program(Q_ref, K_ref, V_ref, D_HEAD_QK**-0.5, bool_mask) - O_ref.backward(dO) - dQ_ref, Q_ref.grad = Q_ref.grad.clone(), None - dK_ref, K_ref.grad = K_ref.grad.clone(), None - dV_ref, V_ref.grad = V_ref.grad.clone(), None - - torch.testing.assert_close(lse, lse_ref, rtol=3e-6, atol=1e-5) - torch.testing.assert_close(O.float(), O_ref, rtol=1e-2, atol=1e-2) - torch.testing.assert_close(dV.float(), dV_ref, rtol=1e-2, atol=1e-2) - torch.testing.assert_close(dK.float(), dK_ref, rtol=1e-2, atol=1e-2) - torch.testing.assert_close(dQ.float(), dQ_ref, rtol=1e-2, atol=1e-2) - print('All checks passed.✅') - - def run(): - Q.grad = K.grad = V.grad = None - flex_attn( - Q, - K, - V, - D_HEAD_QK**-0.5, - mask_fn, - block_mask_fn, - mask_custom_data, - block_mask_custom_data, - )[0].backward(dO) - - with torch.profiler.profile() as prof: - for _ in range(5): - run() - print(prof.key_averages().table(sort_by='cuda_time_total', row_limit=10)) - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--batch', type=int, default=4, help='Batch size') - parser.add_argument('--h', type=int, default=16, help='Number of heads') - parser.add_argument('--n_ctx_qo', type=int, default=4096, help='Context size') - parser.add_argument('--d_head_qk', type=int, default=192, help='Head dimension') - parser.add_argument('--n_ctx_kv', type=int, default=8192, help='Context size') - parser.add_argument('--d_head_vo', type=int, default=128, help='Head dimension') - parser.add_argument('--mask_mode', type=str, default='random', help='Causal flag') - args = parser.parse_args() - main( - args.batch, - args.h, - args.n_ctx_qo, - args.d_head_qk, - args.n_ctx_kv, - args.d_head_vo, - args.mask_mode, - ) diff --git a/src/transform/auto_schedule.cc b/src/transform/auto_schedule.cc index 5f162e2fdb..427c57761d 100644 --- a/src/transform/auto_schedule.cc +++ b/src/transform/auto_schedule.cc @@ -808,6 +808,79 @@ class ScheduleUnitBuilder { seq_body->children = std::move(reordered_children); } + // Reorder & Copy Let-defined variables + auto IsVarDecl = [&] (IRStructure *node) -> bool { + if (!node) return false; + + if (node->IsTask()) { + auto task = static_cast(node); + if (task->stmts.size() == 1) { + return task->stmts[0].as() != nullptr; + } + } + return false; + }; + auto SolveConflictVar = [&] () -> bool { + // LOG(INFO) << "Solving conflict variable..."; + // for (int i = 0; i < n; ++ i) PrintIRStructure(seq_body->children[i].get()); + // LOG(INFO) << "---------------------------"; + for (int i = 0; i < n; ++ i) if (IsVarDecl(seq_body->children[i].get())) { + for (int j = 0; j < n; ++ j) { + if (i == j) continue; + + auto node_i = seq_body->children[i].get(); + auto node_j = seq_body->children[j].get(); + int rem_stage_j = stage_map[node_j]; + + if (!HasDependency(node_i, node_j)) continue; + + LOG(INFO) << "[ScheduleRecursive] Conflict var detection between " << i << " and " << j; + + if (stage_map[node_j] == stage_map[node_i]) continue; + + auto node_i_task = static_cast(node_i); + auto node_i_let_stmt = node_i_task->stmts[0].as(); + + auto iter = ctrl->control->loop_var; + auto step = ctrl->control->step.has_value() ? ctrl->control->step.value() : 1; + auto cloned_value = node_i_let_stmt->value; + auto cloned_let_stmt = LetStmt( + node_i_let_stmt->var.copy_with_suffix(""), + cloned_value, + Evaluate(0) + ); + auto cloned_task = std::make_shared(); + cloned_task->stmts.push_back(cloned_let_stmt); + stage_map[cloned_task.get()] = rem_stage_j; + + for (int k = j; k < n; ++ k) { + auto node_k = seq_body->children[k].get(); + auto task_k = static_cast(node_k); + if (rem_stage_j != stage_map[node_k]) continue; + if (HasDependency(node_i, node_k)) { + for (size_t id = 0; id < task_k->stmts.size(); ++ id) { + task_k->stmts[id] = Substitute( + task_k->stmts[id], + {{node_i_let_stmt->var, cloned_let_stmt->var}} + ); + } + task_k->SubstituteVar(node_i_let_stmt->var, cloned_let_stmt->var); + stage_map[node_k] = rem_stage_j; + } + } + + LOG(INFO) << "Cloned task: " << cloned_task->stmts[0]; + + seq_body->children.insert(seq_body->children.begin() + j, std::move(cloned_task)); + n += 1; + return true; // Conflict resolved, restart the loop + } + } + return false; + }; + int conflict_count = 0; + while (SolveConflictVar() && ++ conflict_count < 100); + for (auto &node : seq_body->children) { auto unit = std::make_shared(); unit->stage = stage_map[node.get()]; @@ -926,16 +999,6 @@ class ScheduleUnitBuilder { if (SameVar(write_var_a, read_var_b)) return true; } - for (const auto &write_var_b : b->GetWriteVars()) { - if (SameVar(write_var_a, write_var_b)) - return true; - } - } - for (const auto &read_var_a : a->GetReadVars()) { - for (const auto &write_var_b : b->GetWriteVars()) { - if (SameVar(read_var_a, write_var_b)) - return true; - } } return false; } @@ -990,22 +1053,6 @@ class ScheduleUnitBuilder { return true; } } - for (const auto &write_var_a : a->GetWriteVars()) { - for (const auto &read_var_b : b->GetReadVars()) { - if (SameVar(write_var_a, read_var_b)) - return true; - } - for (const auto &write_var_b : b->GetWriteVars()) { - if (SameVar(write_var_a, write_var_b)) - return true; - } - } - for (const auto &read_var_a : a->GetReadVars()) { - for (const auto &write_var_b : b->GetWriteVars()) { - if (SameVar(read_var_a, write_var_b)) - return true; - } - } return false; } @@ -1316,6 +1363,7 @@ class IRStructureBuilder : public StmtVisitor { task_node->stmts.push_back(GetRef(op)); AnalyzeMemoryExpr(op->condition, task_node.get()); + AnalyzeResourceUsage(Evaluate(op->condition), task_node.get(), true); // Analyze both branches for resource usage AnalyzeResourceUsage(op->then_case, task_node.get()); @@ -1390,7 +1438,7 @@ class IRStructureBuilder : public StmtVisitor { int64_t thread_count_ = 1; Target target_; - void AnalyzeResourceUsage(const Stmt &stmt, TaskNode *task_node) { + void AnalyzeResourceUsage(const Stmt &stmt, TaskNode *task_node, bool only_variables = false) { // Recursively analyze statements to determine resource usage struct ResourceAnalyzer : public StmtExprVisitor { TaskNode *task_node; @@ -1440,7 +1488,8 @@ class IRStructureBuilder : public StmtVisitor { // Check if this is a TMA copy operation if (op->op.same_as(copy_op)) { - bool found_global = false; + bool found_global = false, found_shared = false; + int idx_global = -1, idx_shared = -1; for (unsigned idx = 0; idx != 2; ++idx) { auto region = Downcast(op->args[idx]); if (const auto *buffer_load = @@ -1450,13 +1499,21 @@ class IRStructureBuilder : public StmtVisitor { MemoryType mem_type = GetMemoryTypeFromScope(scope); if (mem_type == MemoryType::kGlobal) { found_global = true; - if (idx == 0) { - found_tma_load = true; - } + idx_global = idx; + } + if (mem_type == MemoryType::kShared) { + found_shared = true; + idx_shared = idx; } } } - found_tma |= found_global; + found_tma = false; + if (found_global && found_shared) { + if (idx_global == 0 && idx_shared == 1) + found_tma = true; + if (idx_global == 1 && idx_shared == 0) + found_tma = true; + } } else if (op->op.same_as(gemm_py_op) || op->op.same_as(gemm_op)) { found_tensor = true; @@ -1532,31 +1589,33 @@ class IRStructureBuilder : public StmtVisitor { ResourceAnalyzer analyzer(task_node, target_, thread_count_); analyzer(stmt); - // Set task node flags based on what was found - if (analyzer.found_tma) { - task_node->SetUsesTMACore(true); - if (analyzer.found_tma_load) { - task_node->SetHasTMALoad(true); + if (!only_variables) { + // Set task node flags based on what was found + if (analyzer.found_tma) { + task_node->SetUsesTMACore(true); + if (analyzer.found_tma_load) { + task_node->SetHasTMALoad(true); + } } - } - if (analyzer.found_tensor) { - task_node->SetUsesTensorCore(true); - // Set Tensor Core shape information if available - for (const auto &shape : analyzer.tensor_core_shapes) { - if (shape.m > 0 && shape.n > 0 && shape.k > 0) { - task_node->AddTensorCoreShape(shape.m, shape.n, shape.k); + if (analyzer.found_tensor) { + task_node->SetUsesTensorCore(true); + // Set Tensor Core shape information if available + for (const auto &shape : analyzer.tensor_core_shapes) { + if (shape.m > 0 && shape.n > 0 && shape.k > 0) { + task_node->AddTensorCoreShape(shape.m, shape.n, shape.k); + } + } + // Set GemmInst information + if (analyzer.has_gemm_inst) { + task_node->SetGemmInst(analyzer.gemm_inst); } } - // Set GemmInst information - if (analyzer.has_gemm_inst) { - task_node->SetGemmInst(analyzer.gemm_inst); + // If neither TMA nor Tensor core was used, and CUDA operations were found, + // set CUDA core flag + if (!analyzer.found_tma && !analyzer.found_tensor) { + task_node->SetUsesCUDACore(true); } } - // If neither TMA nor Tensor core was used, and CUDA operations were found, - // set CUDA core flag - if (!analyzer.found_tma && !analyzer.found_tensor) { - task_node->SetUsesCUDACore(true); - } // Analyze memory access regions MemoryAccessDetector memory_detector; @@ -1673,7 +1732,7 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) { ICHECK(ir_structure) << "IRStructure is null (empty body?)"; // First print the summary view - // PrintIRStructure(ir_structure.get()); + PrintIRStructure(ir_structure.get()); // Then print all statements // PrintAllStmts(ir_structure.get()); @@ -1696,9 +1755,6 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) { unit_builder.SetEnableWarpPartition(config.enable_warp_partition); bool double_thread = unit_builder.Build(ir_structure.get()); - PrintIRStructure(ir_structure.get()); - LOG(FATAL) << "Early exit for debugging purposes\n"; - if (!config.enable_warpgroup_partition) { Stmt new_body = ConvertIRStructureToStmt(ir_structure.get(), enable_epi); @@ -1716,7 +1772,7 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) { } // Print the modified summary view - PrintIRStructure(ir_structure.get()); + // PrintIRStructure(ir_structure.get()); // Analyze buffer dependencies and insert barriers before warpgroup // partition @@ -1833,54 +1889,303 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) { return CreatePrimFuncPass(pass_func, 0, "tl.AutoSchedule", {}); } -// Helper function to clone IRStructure with warpgroup filter +// Helper: check if a TaskNode is a LetDecl (single LetStmt with empty body) +static bool IsLetDeclTask(const TaskNode *task) { + return task->stmts.size() == 1 && + task->stmts[0].as() != nullptr; +} + +// Helper: check if an IRStructure node is a LetDecl task (or a ScheduleUnit +// wrapping one) +static bool IsLetDeclNode(const IRStructure *node) { + if (!node) return false; + if (node->IsTask()) { + return IsLetDeclTask(static_cast(node)); + } + if (node->IsScheduleUnit()) { + auto unit = static_cast(node); + return unit->child && unit->child->IsTask() && + IsLetDeclTask(static_cast(unit->child.get())); + } + return false; +} + +// Helper: check if an IRStructure subtree contains any LetDecl tasks +static bool ContainsLetDecl(const IRStructure *node) { + if (!node) return false; + if (IsLetDeclNode(node)) return true; + if (node->IsSequence()) { + auto seq = static_cast(node); + for (const auto &child : seq->children) { + if (ContainsLetDecl(child.get())) return true; + } + } else if (node->IsControl()) { + auto ctrl = static_cast(node); + return ContainsLetDecl(ctrl->child.get()); + } else if (node->IsWrapper()) { + auto wrapper = static_cast(node); + return ContainsLetDecl(wrapper->child.get()); + } else if (node->IsScheduleUnit()) { + auto unit = static_cast(node); + return ContainsLetDecl(unit->child.get()); + } + return false; +} + +// Helper function to clone IRStructure with warpgroup filter. std::shared_ptr -CloneIRStructureWithWarpgroupFilter(IRStructure *node, int warpgroup_id) { - if (!node || !node->containWarpgroupId(warpgroup_id)) - return nullptr; +CloneIRStructureWithWarpgroupFilter(IRStructure *node, int warpgroup_id, + Map &var_remap) { + if (!node) return nullptr; if (node->IsTask()) { auto task = static_cast(node); - return task->Clone(); + + // LetDecl tasks are always included in every warp group clone. + // Create a fresh variable copy so the two warp groups use different names. + if (IsLetDeclTask(task)) { + const auto *let = task->stmts[0].as(); + auto new_var = let->var.copy_with_suffix(""); + // Substitute previously renamed variables in the value expression. + PrimExpr new_value = var_remap.empty() + ? let->value + : Substitute(let->value, var_remap); + var_remap.Set(let->var, new_var); + auto new_task = std::make_shared(); + new_task->stmts.push_back(LetStmt(new_var, new_value, Evaluate(0))); + return new_task; + } + + // Non-LetDecl tasks: only include if warp group matches + if (!node->containWarpgroupId(warpgroup_id)) return nullptr; + auto cloned = task->Clone(); + // Substitute renamed LetDecl variables in task statements + if (!var_remap.empty()) { + auto ct = static_cast(cloned.get()); + for (size_t i = 0; i < ct->stmts.size(); ++i) { + ct->stmts[i] = Substitute(ct->stmts[i], var_remap); + } + } + return cloned; } else if (node->IsSequence()) { + // A SequenceNode is included if it contains the target warp group + // OR if it contains LetDecl tasks (which are always needed). + if (!node->containWarpgroupId(warpgroup_id) && !ContainsLetDecl(node)) + return nullptr; auto seq = static_cast(node); auto new_seq = std::make_shared(); for (const auto &child : seq->children) { - auto new_child = - CloneIRStructureWithWarpgroupFilter(child.get(), warpgroup_id); + auto new_child = CloneIRStructureWithWarpgroupFilter( + child.get(), warpgroup_id, var_remap); if (new_child) { new_seq->children.push_back(std::move(new_child)); } } + if (new_seq->children.empty()) return nullptr; return new_seq; } else if (node->IsControl()) { + // A ControlNode is included if it contains the target warp group + // OR if it contains LetDecl tasks. + if (!node->containWarpgroupId(warpgroup_id) && !ContainsLetDecl(node)) + return nullptr; auto ctrl = static_cast(node); auto new_ctrl = std::make_shared(); new_ctrl->control = ctrl->control; new_ctrl->SetPromote(ctrl->hasPromote()); - new_ctrl->child = - CloneIRStructureWithWarpgroupFilter(ctrl->child.get(), warpgroup_id); + new_ctrl->child = CloneIRStructureWithWarpgroupFilter( + ctrl->child.get(), warpgroup_id, var_remap); return new_ctrl; } else if (node->IsWrapper()) { + if (!node->containWarpgroupId(warpgroup_id) && !ContainsLetDecl(node)) + return nullptr; auto wrapper = static_cast(node); auto new_wrapper = std::make_shared(); + // Keep the wrapper statement as-is (do NOT rename LetStmt wrappers here; + // only LetDecl TaskNodes get renamed). new_wrapper->wrapper = wrapper->wrapper; - new_wrapper->child = - CloneIRStructureWithWarpgroupFilter(wrapper->child.get(), warpgroup_id); + new_wrapper->child = CloneIRStructureWithWarpgroupFilter( + wrapper->child.get(), warpgroup_id, var_remap); return new_wrapper; } else if (node->IsScheduleUnit()) { auto unit = static_cast(node); + bool child_is_let_decl = IsLetDeclNode(unit->child.get()); + + // Include the ScheduleUnit if the child is a LetDecl or the warp group + // matches. + if (!child_is_let_decl && !node->containWarpgroupId(warpgroup_id)) + return nullptr; + auto new_unit = std::make_shared(); - new_unit->before[warpgroup_id] = unit->before[warpgroup_id]; - new_unit->after[warpgroup_id] = unit->after[warpgroup_id]; new_unit->stage = unit->stage; - new_unit->child = - CloneIRStructureWithWarpgroupFilter(unit->child.get(), warpgroup_id); + new_unit->child = CloneIRStructureWithWarpgroupFilter( + unit->child.get(), warpgroup_id, var_remap); + + if (!child_is_let_decl) { + // Copy before/after for the target warp group + new_unit->before[warpgroup_id] = unit->before[warpgroup_id]; + new_unit->after[warpgroup_id] = unit->after[warpgroup_id]; + // Substitute renamed LetDecl variables in before/after stmts + if (!var_remap.empty()) { + for (auto &s : new_unit->before[warpgroup_id]) { + s = Substitute(s, var_remap); + } + for (auto &s : new_unit->after[warpgroup_id]) { + s = Substitute(s, var_remap); + } + } + } return new_unit; } LOG(FATAL); } +// Entry point overload — creates a fresh var_remap per call +std::shared_ptr +CloneIRStructureWithWarpgroupFilter(IRStructure *node, int warpgroup_id) { + Map var_remap; + return CloneIRStructureWithWarpgroupFilter(node, warpgroup_id, var_remap); +} + +// Simple visitor to collect all Var references from statements/expressions +class VarRefCollector : public StmtExprVisitor { +public: + std::unordered_set vars; + void VisitExpr_(const VarNode *op) override { vars.insert(op); } +}; + +// Remove LetDecl definitions whose variables are not referenced by any +// non-LetDecl task in the IR tree. After warp-group partitioning some +// LetDecl values may access out-of-bounds indices because the consumer +// tasks that used them ended up in the other warp group. +std::shared_ptr +RemoveUnusedLetDecls(std::shared_ptr root) { + if (!root) return nullptr; + + // Phase 1: Collect LetDecl definitions and variable references from + // non-LetDecl nodes (task stmts and ScheduleUnit before/after). + struct LetDeclEntry { + const VarNode *var; + PrimExpr value; + }; + std::vector let_decls; + std::unordered_set referenced_vars; + + std::function collect = + [&](const IRStructure *node) { + if (!node) return; + if (node->IsTask()) { + auto task = static_cast(node); + if (IsLetDeclTask(task)) { + const auto *let = task->stmts[0].as(); + let_decls.push_back({let->var.get(), let->value}); + } else { + VarRefCollector collector; + for (const auto &stmt : task->stmts) { + collector(stmt); + } + referenced_vars.insert(collector.vars.begin(), + collector.vars.end()); + } + } else if (node->IsSequence()) { + for (const auto &child : + static_cast(node)->children) { + collect(child.get()); + } + } else if (node->IsControl()) { + collect(static_cast(node)->child.get()); + } else if (node->IsWrapper()) { + collect(static_cast(node)->child.get()); + } else if (node->IsScheduleUnit()) { + auto unit = static_cast(node); + collect(unit->child.get()); + VarRefCollector collector; + for (const auto &stmts : unit->before) { + for (const auto &s : stmts) collector(s); + } + for (const auto &stmts : unit->after) { + for (const auto &s : stmts) collector(s); + } + referenced_vars.insert(collector.vars.begin(), + collector.vars.end()); + } + }; + collect(root.get()); + + // Phase 2: Transitive closure — if a LetDecl var is referenced, + // all vars in its value expression are transitively referenced too. + bool changed = true; + while (changed) { + changed = false; + for (const auto &entry : let_decls) { + if (referenced_vars.count(entry.var)) { + VarRefCollector collector; + collector(entry.value); + for (const auto *v : collector.vars) { + if (!referenced_vars.count(v)) { + referenced_vars.insert(v); + changed = true; + } + } + } + } + } + + // Phase 3: Filter the tree — remove LetDecl tasks for unused vars. + std::function( + const std::shared_ptr &)> + filter_tree = [&](const std::shared_ptr &node) + -> std::shared_ptr { + if (!node) return nullptr; + if (node->IsTask()) { + if (IsLetDeclTask(static_cast(node.get()))) { + const auto *let = + static_cast(node.get()) + ->stmts[0] + .as(); + if (!referenced_vars.count(let->var.get())) { + return nullptr; // Remove unused LetDecl + } + } + return node; + } else if (node->IsSequence()) { + auto seq = static_cast(node.get()); + auto new_seq = std::make_shared(); + for (const auto &child : seq->children) { + auto filtered = filter_tree(child); + if (filtered) new_seq->children.push_back(std::move(filtered)); + } + if (new_seq->children.empty()) return nullptr; + return new_seq; + } else if (node->IsControl()) { + auto ctrl = static_cast(node.get()); + auto new_ctrl = std::make_shared(); + new_ctrl->control = ctrl->control; + new_ctrl->SetPromote(ctrl->hasPromote()); + new_ctrl->child = filter_tree(ctrl->child); + if (!new_ctrl->child) return nullptr; + return new_ctrl; + } else if (node->IsWrapper()) { + auto wrapper = static_cast(node.get()); + auto new_wrapper = std::make_shared(); + new_wrapper->wrapper = wrapper->wrapper; + new_wrapper->child = filter_tree(wrapper->child); + return new_wrapper; + } else if (node->IsScheduleUnit()) { + auto unit = static_cast(node.get()); + auto new_unit = std::make_shared(); + new_unit->stage = unit->stage; + new_unit->before = unit->before; + new_unit->after = unit->after; + new_unit->child = filter_tree(unit->child); + if (!new_unit->child) return nullptr; + return new_unit; + } + return node; + }; + + return filter_tree(root); +} + class SimtCopyDetector : public StmtExprVisitor { public: static bool Detect(const Stmt &stmt) { @@ -2083,19 +2388,26 @@ Stmt ConvertIRStructureToStmt(IRStructure *root, const bool outer_enable_epi) { stmts.push_back(stmt); } } - Map substitution; - PrimExpr condition = - And(loop_var < loop_extent, loop_var >= loop_start); - if (unit->stage == min_stages) { - condition = loop_var >= loop_start; - } - if (unit->stage == max_stages) { - condition = loop_var < loop_extent; - } - Stmt stmt = IfThenElse(condition, SeqStmt::Flatten(stmts)); + Map substitution, substitution_cond; substitution.Set(loop_var, loop_var - loop_step * (max_stages - unit->stage)); - steady.push_back(Substitute(stmt, substitution)); + substitution_cond.Set(loop_var, + Max(loop_start, Min(loop_start + loop_extent - loop_step, loop_var - loop_step * (max_stages - unit->stage)))); + if (IsLetDeclNode(unit->child.get())) { + Stmt stmt = SeqStmt::Flatten(stmts); + steady.push_back(Substitute(stmt, substitution_cond)); + } else { + PrimExpr condition = + And(loop_var < loop_extent, loop_var >= loop_start); + if (unit->stage == min_stages) { + condition = loop_var >= loop_start; + } + if (unit->stage == max_stages) { + condition = loop_var < loop_extent; + } + Stmt stmt = IfThenElse(condition, SeqStmt::Flatten(stmts)); + steady.push_back(Substitute(stmt, substitution)); + } } Stmt new_body = SeqStmt::Flatten(steady); auto new_var = loop_var.copy_with_suffix(""); @@ -2391,19 +2703,26 @@ Stmt ApplyWarpgroupPartitionToIRStructure( stmts.push_back(stmt); } } - Map substitution; - PrimExpr condition = - And(loop_var < loop_extent, loop_var >= loop_start); - if (unit->stage == min_stages) { - condition = loop_var >= loop_start; - } - if (unit->stage == max_stages) { - condition = loop_var < loop_extent; - } - Stmt stmt = IfThenElse(condition, SeqStmt::Flatten(stmts)); + Map substitution, substitution_cond; substitution.Set(loop_var, loop_var - loop_step * (max_stages - unit->stage)); - steady.push_back(Substitute(stmt, substitution)); + substitution_cond.Set(loop_var, + Max(loop_start, Min(loop_start + loop_extent - loop_step, loop_var - loop_step * (max_stages - unit->stage)))); + if (IsLetDeclNode(unit->child.get())) { + Stmt stmt = SeqStmt::Flatten(stmts); + steady.push_back(Substitute(stmt, substitution_cond)); + } else { + PrimExpr condition = + And(loop_var < loop_extent, loop_var >= loop_start); + if (unit->stage == min_stages) { + condition = loop_var >= loop_start; + } + if (unit->stage == max_stages) { + condition = loop_var < loop_extent; + } + Stmt stmt = IfThenElse(condition, SeqStmt::Flatten(stmts)); + steady.push_back(Substitute(stmt, substitution)); + } } Stmt new_body = SeqStmt::Flatten(steady); auto new_var = loop_var.copy_with_suffix(""); @@ -2626,8 +2945,10 @@ Stmt ApplyWarpgroupPartitionToIRStructure( root, is_epi_top_level_index, -1) : nullptr; - auto wg0_structure = CloneIRStructureWithWarpgroupFilter(root, 0); - auto wg1_structure = CloneIRStructureWithWarpgroupFilter(root, 1); + auto wg0_structure = + RemoveUnusedLetDecls(CloneIRStructureWithWarpgroupFilter(root, 0)); + auto wg1_structure = + RemoveUnusedLetDecls(CloneIRStructureWithWarpgroupFilter(root, 1)); bool wg_pro_neutral_has_stmts = wg_pro_neutral_structure diff --git a/src/transform/auto_schedule/ir_structure.h b/src/transform/auto_schedule/ir_structure.h index 715add0a89..5682fe0582 100644 --- a/src/transform/auto_schedule/ir_structure.h +++ b/src/transform/auto_schedule/ir_structure.h @@ -160,7 +160,7 @@ class TaskNode : public IRStructure { // Latency estimation int64_t GetLatency() const override { return latency_; } int64_t GetII() const override { return ii_; } - + // Setters void SetUsesCUDACore(bool value) override { uses_cuda_core_ = value; } void SetUsesTMACore(bool value) override { uses_tma_core_ = value; } @@ -171,6 +171,24 @@ class TaskNode : public IRStructure { void SetWriteRegions(const std::vector ®ions) override { write_regions_ = regions; } + void SetReadVars(const std::vector &vars) { + read_vars_ = vars; + } + void SetWriteVars(const std::vector &vars) { + write_vars_ = vars; + } + void SubstituteVar(const Var &old_var, const Var &new_var) { + for (auto &var : read_vars_) { + if (var.same_as(old_var)) { + var = new_var; + } + } + for (auto &var : write_vars_) { + if (var.same_as(old_var)) { + var = new_var; + } + } + } void SetLatency(int64_t latency) override { latency_ = latency; } void SetII(int64_t ii) override { ii_ = ii; } @@ -954,6 +972,12 @@ inline void PrintIRStructure(const IRStructure *node, int indent = 0) { for (auto ®ion : task->GetWriteRegions()) { LOG(INFO) << indent_str << " Write Region: " << region; } + for (auto &var : task->GetReadVars()) { + LOG(INFO) << indent_str << " Read Var: " << var; + } + for (auto &var : task->GetWriteVars()) { + LOG(INFO) << indent_str << " Write Var: " << var; + } LOG(INFO) << indent_str << " uses_cuda_core: " << task->UsesCUDACore(); LOG(INFO) << indent_str << " uses_tma_core: " << task->UsesTMACore(); LOG(INFO) << indent_str << " uses_tensor_core: " << task->UsesTensorCore(); diff --git a/src/transform/auto_schedule/memory_detector.h b/src/transform/auto_schedule/memory_detector.h index dbf1136ebc..0364b2af20 100644 --- a/src/transform/auto_schedule/memory_detector.h +++ b/src/transform/auto_schedule/memory_detector.h @@ -40,7 +40,6 @@ class MemoryAccessDetector : public StmtExprVisitor { let_bindings_.clear(); read_vars_.clear(); write_vars_.clear(); - let_defined_vars_.clear(); operator()(stmt); } @@ -55,7 +54,6 @@ class MemoryAccessDetector : public StmtExprVisitor { let_bindings_.clear(); read_vars_.clear(); write_vars_.clear(); - let_defined_vars_.clear(); operator()(expr); } @@ -98,8 +96,6 @@ class MemoryAccessDetector : public StmtExprVisitor { std::vector read_vars_; /*! \brief The set of variables that are written in the current block. */ std::vector write_vars_; - /*! \brief The set of variables defined by let bindings in the current scope. */ - std::set let_defined_vars_; /*! * \brief Update read/write buffers and regions with provided buffer and @@ -267,16 +263,13 @@ class MemoryAccessDetector : public StmtExprVisitor { void VisitStmt_(const LetStmtNode *op) override { let_bindings_[op->var.get()] = op->value; - let_defined_vars_.insert(op->var.get()); UpdateWriteVar(op->var); StmtExprVisitor::VisitStmt(op->body); let_bindings_.erase(op->var.get()); } void VisitExpr_(const VarNode *op) override { - if (let_defined_vars_.count(op)) { - UpdateReadVar(tvm::ffi::GetRef(op)); - } + UpdateReadVar(tvm::ffi::GetRef(op)); StmtExprVisitor::VisitExpr_(op); } diff --git a/tilelang/engine/phase.py b/tilelang/engine/phase.py index b1b3e65624..66e1f1b08f 100644 --- a/tilelang/engine/phase.py +++ b/tilelang/engine/phase.py @@ -333,6 +333,6 @@ def OptimizeForTarget(mod: IRModule, target: Target) -> IRModule: # Transform threadblock to persistent threadblock mod = tilelang.transform.PersistThreadblock()(mod) - print(mod) + # print(mod) return mod From 56385cb14b9f68448128464aad364143063c6a51 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Fri, 27 Mar 2026 18:01:02 +0800 Subject: [PATCH 3/6] Remove debug outputs --- src/transform/auto_schedule.cc | 6 +++--- tilelang/engine/phase.py | 3 --- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/src/transform/auto_schedule.cc b/src/transform/auto_schedule.cc index 427c57761d..d3443c93d0 100644 --- a/src/transform/auto_schedule.cc +++ b/src/transform/auto_schedule.cc @@ -834,7 +834,7 @@ class ScheduleUnitBuilder { if (!HasDependency(node_i, node_j)) continue; - LOG(INFO) << "[ScheduleRecursive] Conflict var detection between " << i << " and " << j; + // LOG(INFO) << "[ScheduleRecursive] Conflict var detection between " << i << " and " << j; if (stage_map[node_j] == stage_map[node_i]) continue; @@ -869,7 +869,7 @@ class ScheduleUnitBuilder { } } - LOG(INFO) << "Cloned task: " << cloned_task->stmts[0]; + // LOG(INFO) << "Cloned task: " << cloned_task->stmts[0]; seq_body->children.insert(seq_body->children.begin() + j, std::move(cloned_task)); n += 1; @@ -1732,7 +1732,7 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) { ICHECK(ir_structure) << "IRStructure is null (empty body?)"; // First print the summary view - PrintIRStructure(ir_structure.get()); + // PrintIRStructure(ir_structure.get()); // Then print all statements // PrintAllStmts(ir_structure.get()); diff --git a/tilelang/engine/phase.py b/tilelang/engine/phase.py index 66e1f1b08f..c1072db995 100644 --- a/tilelang/engine/phase.py +++ b/tilelang/engine/phase.py @@ -201,7 +201,6 @@ def LowerAndLegalize(mod: IRModule, target: Target) -> IRModule: # print(mod) mod = tilelang.transform.AutoSchedule(False)(mod) mod = tilelang.transform.Simplify()(mod) - print(mod) # Set layouts for reducers mod = tilelang.transform.LayoutReducer()(mod) # Infer memory layouts for fragments and shared memory @@ -333,6 +332,4 @@ def OptimizeForTarget(mod: IRModule, target: Target) -> IRModule: # Transform threadblock to persistent threadblock mod = tilelang.transform.PersistThreadblock()(mod) - # print(mod) - return mod From 3bea503c08f84910c68a6538fa8bc4f948de3915 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Mon, 30 Mar 2026 15:42:32 +0800 Subject: [PATCH 4/6] Clean dead code --- src/transform/auto_schedule/ir_structure.cc | 12 ---- src/transform/auto_schedule/ir_structure.h | 79 ++------------------- 2 files changed, 4 insertions(+), 87 deletions(-) diff --git a/src/transform/auto_schedule/ir_structure.cc b/src/transform/auto_schedule/ir_structure.cc index 45531a6e11..b23330e126 100644 --- a/src/transform/auto_schedule/ir_structure.cc +++ b/src/transform/auto_schedule/ir_structure.cc @@ -208,18 +208,6 @@ void SequenceNode::SetLatency(int64_t latency) { latency_ = latency; } void SequenceNode::SetII(int64_t ii) { ii_ = ii; } -void SequenceNode::AddReadRegion(const BufferRegion ®ion) { - if (!children.empty() && children[0]) { - children[0]->AddReadRegion(region); - } -} - -void SequenceNode::AddWriteRegion(const BufferRegion ®ion) { - if (!children.empty() && children[0]) { - children[0]->AddWriteRegion(region); - } -} - std::shared_ptr SequenceNode::Clone() const { auto new_seq = std::make_shared(); new_seq->children.reserve(children.size()); diff --git a/src/transform/auto_schedule/ir_structure.h b/src/transform/auto_schedule/ir_structure.h index 5682fe0582..5c0b2fffd6 100644 --- a/src/transform/auto_schedule/ir_structure.h +++ b/src/transform/auto_schedule/ir_structure.h @@ -101,13 +101,6 @@ class IRStructure { virtual void SetLatency(int64_t latency) = 0; virtual void SetII(int64_t ii) = 0; - // Helper methods to add regions (for incremental analysis) - virtual void AddReadRegion(const BufferRegion ®ion) = 0; - virtual void AddWriteRegion(const BufferRegion ®ion) = 0; - - // Helper methods to add variables - virtual void AddReadVar(const Var &var) = 0; - virtual void AddWriteVar(const Var &var) = 0; // Recursive region collection method virtual void CollectRegions( @@ -281,7 +274,7 @@ class TaskNode : public IRStructure { std::shared_ptr Clone() const override; // Helper methods to add regions (for incremental analysis) - void AddReadRegion(const BufferRegion ®ion) override { + void AddReadRegion(const BufferRegion ®ion) { // Check for duplicate regions for (const auto &existing : read_regions_) { if (existing->buffer.same_as(region->buffer) && @@ -292,7 +285,7 @@ class TaskNode : public IRStructure { read_regions_.push_back(region); } - void AddWriteRegion(const BufferRegion ®ion) override { + void AddWriteRegion(const BufferRegion ®ion) { // Check for duplicate regions for (const auto &existing : write_regions_) { if (existing->buffer.same_as(region->buffer) && @@ -303,8 +296,8 @@ class TaskNode : public IRStructure { write_regions_.push_back(region); } - void AddReadVar(const Var &var) override { read_vars_.push_back(var); } - void AddWriteVar(const Var &var) override { write_vars_.push_back(var); } + void AddReadVar(const Var &var) { read_vars_.push_back(var); } + void AddWriteVar(const Var &var) { write_vars_.push_back(var); } void CollectRegions( std::vector &result, @@ -413,23 +406,6 @@ class ControlNode : public IRStructure { void SetLatency(int64_t latency) override { latency_ = latency; } void SetII(int64_t ii) override { ii_ = ii; } - // Helper methods to add regions (delegate to child) - void AddReadRegion(const BufferRegion ®ion) override { - if (child) - child->AddReadRegion(region); - } - void AddWriteRegion(const BufferRegion ®ion) override { - if (child) - child->AddWriteRegion(region); - } - - void AddReadVar(const Var &var) override { - // ignore - } - void AddWriteVar(const Var &var) override { - // ignore - } - void CollectRegions( std::vector &result, std::set>> &visited) const override; @@ -517,23 +493,6 @@ class WrapperNode : public IRStructure { void SetLatency(int64_t latency) override { latency_ = latency; } void SetII(int64_t ii) override { ii_ = ii; } - // Helper methods to add regions (delegate to child) - void AddReadRegion(const BufferRegion ®ion) override { - if (child) - child->AddReadRegion(region); - } - void AddWriteRegion(const BufferRegion ®ion) override { - if (child) - child->AddWriteRegion(region); - } - - void AddReadVar(const Var &var) override { - // ignore - } - void AddWriteVar(const Var &var) override { - // ignore - } - void CollectRegions( std::vector &result, std::set>> &visited) const override; @@ -620,25 +579,6 @@ class ScheduleUnit : public IRStructure { void SetLatency(int64_t latency) override { latency_ = latency; } void SetII(int64_t ii) override { ii_ = ii; } - // Helper methods to add regions (delegate to child) - void AddReadRegion(const BufferRegion ®ion) override { - if (child) - child->AddReadRegion(region); - } - void AddWriteRegion(const BufferRegion ®ion) override { - if (child) - child->AddWriteRegion(region); - } - - void AddReadVar(const Var &var) override { - if (child) - child->AddReadVar(var); - } - void AddWriteVar(const Var &var) override { - if (child) - child->AddWriteVar(var); - } - void CollectRegions( std::vector &result, std::set>> &visited) const override; @@ -696,17 +636,6 @@ class SequenceNode : public IRStructure { void SetLatency(int64_t latency) override; void SetII(int64_t ii) override; - // Helper methods to add regions (delegate to first child if exists) - void AddReadRegion(const BufferRegion ®ion) override; - void AddWriteRegion(const BufferRegion ®ion) override; - - void AddReadVar(const Var &var) override { - // ignore - } - void AddWriteVar(const Var &var) override { - // ignore - } - void CollectRegions( std::vector &result, std::set>> &visited) const override; From 551945d5d842ac703953c813d8cf0c5cf6725b65 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Mon, 30 Mar 2026 17:30:10 +0800 Subject: [PATCH 5/6] Remove debug outputs --- src/transform/auto_schedule.cc | 6 +----- tilelang/engine/phase.py | 2 -- 2 files changed, 1 insertion(+), 7 deletions(-) diff --git a/src/transform/auto_schedule.cc b/src/transform/auto_schedule.cc index d3443c93d0..2fb0b0543e 100644 --- a/src/transform/auto_schedule.cc +++ b/src/transform/auto_schedule.cc @@ -821,9 +821,6 @@ class ScheduleUnitBuilder { return false; }; auto SolveConflictVar = [&] () -> bool { - // LOG(INFO) << "Solving conflict variable..."; - // for (int i = 0; i < n; ++ i) PrintIRStructure(seq_body->children[i].get()); - // LOG(INFO) << "---------------------------"; for (int i = 0; i < n; ++ i) if (IsVarDecl(seq_body->children[i].get())) { for (int j = 0; j < n; ++ j) { if (i == j) continue; @@ -869,8 +866,6 @@ class ScheduleUnitBuilder { } } - // LOG(INFO) << "Cloned task: " << cloned_task->stmts[0]; - seq_body->children.insert(seq_body->children.begin() + j, std::move(cloned_task)); n += 1; return true; // Conflict resolved, restart the loop @@ -878,6 +873,7 @@ class ScheduleUnitBuilder { } return false; }; + // Resolve conflicts until no more conflicts exist or max iterations reached (to avoid infinite loop) int conflict_count = 0; while (SolveConflictVar() && ++ conflict_count < 100); diff --git a/tilelang/engine/phase.py b/tilelang/engine/phase.py index 01ab2d9ad8..9ce8b28187 100644 --- a/tilelang/engine/phase.py +++ b/tilelang/engine/phase.py @@ -183,8 +183,6 @@ def LowerAndLegalize(mod: IRModule, target: Target) -> IRModule: if allow_autoschedule(): # Auto schedule for high-level operations mod = tilelang.transform.IfConditionExtract()(mod) - # print("IfConditionExtract done") - # print(mod) mod = tilelang.transform.AutoSchedule(False)(mod) mod = tilelang.transform.Simplify()(mod) # Set layouts for reducers From 0687a3fe3b2cef220bcda9ebefb83f11178bc9f8 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Mon, 30 Mar 2026 18:06:54 +0800 Subject: [PATCH 6/6] Format code --- src/transform/auto_schedule.cc | 249 ++++++++++-------- src/transform/auto_schedule/ir_structure.cc | 6 +- src/transform/auto_schedule/ir_structure.h | 26 +- src/transform/auto_schedule/memory_detector.h | 6 +- src/transform/if_condition_extract.cc | 20 +- tilelang/transform/__init__.py | 2 + 6 files changed, 177 insertions(+), 132 deletions(-) diff --git a/src/transform/auto_schedule.cc b/src/transform/auto_schedule.cc index 2fb0b0543e..3b25a8222b 100644 --- a/src/transform/auto_schedule.cc +++ b/src/transform/auto_schedule.cc @@ -70,10 +70,11 @@ using namespace tir; using ffi::GetRef; // Extract all sequencial task nodes from the IR structure tree -void GatherTaskNodesSingle(const std::shared_ptr &node, - std::vector> &task_nodes); +void GatherTaskNodesSingle( + const std::shared_ptr &node, + std::vector> &task_nodes); void GatherTaskNodes(const std::vector> &nodes, - std::vector> &task_nodes) { + std::vector> &task_nodes) { for (const auto &node : nodes) { if (node->IsTask()) { task_nodes.emplace_back(node); @@ -82,8 +83,10 @@ void GatherTaskNodes(const std::vector> &nodes, GatherTaskNodes(seq->children, task_nodes); } else if (node->IsWrapper()) { auto wrapper = static_cast(node.get()); - if (wrapper->task) task_nodes.emplace_back(wrapper->task); - if (wrapper->child) GatherTaskNodesSingle(wrapper->child, task_nodes); + if (wrapper->task) + task_nodes.emplace_back(wrapper->task); + if (wrapper->child) + GatherTaskNodesSingle(wrapper->child, task_nodes); } else if (node->IsControl()) { task_nodes.emplace_back(node); } else { @@ -92,8 +95,9 @@ void GatherTaskNodes(const std::vector> &nodes, } } -void GatherTaskNodesSingle(const std::shared_ptr &node, - std::vector> &task_nodes) { +void GatherTaskNodesSingle( + const std::shared_ptr &node, + std::vector> &task_nodes) { return GatherTaskNodes({node}, task_nodes); } @@ -605,7 +609,8 @@ class ScheduleUnitBuilder { // with distance-aware dependencies void Z3SchedulePythonLoop(ControlNode *ctrl) { if (ctrl->child == nullptr) { - LOG(WARNING) << "Z3SchedulePythonLoop called on a control node without child"; + LOG(WARNING) + << "Z3SchedulePythonLoop called on a control node without child"; return; } std::vector> flat_children; @@ -809,8 +814,9 @@ class ScheduleUnitBuilder { } // Reorder & Copy Let-defined variables - auto IsVarDecl = [&] (IRStructure *node) -> bool { - if (!node) return false; + auto IsVarDecl = [&](IRStructure *node) -> bool { + if (!node) + return false; if (node->IsTask()) { auto task = static_cast(node); @@ -820,62 +826,71 @@ class ScheduleUnitBuilder { } return false; }; - auto SolveConflictVar = [&] () -> bool { - for (int i = 0; i < n; ++ i) if (IsVarDecl(seq_body->children[i].get())) { - for (int j = 0; j < n; ++ j) { - if (i == j) continue; - - auto node_i = seq_body->children[i].get(); - auto node_j = seq_body->children[j].get(); - int rem_stage_j = stage_map[node_j]; - - if (!HasDependency(node_i, node_j)) continue; - - // LOG(INFO) << "[ScheduleRecursive] Conflict var detection between " << i << " and " << j; - - if (stage_map[node_j] == stage_map[node_i]) continue; - - auto node_i_task = static_cast(node_i); - auto node_i_let_stmt = node_i_task->stmts[0].as(); - - auto iter = ctrl->control->loop_var; - auto step = ctrl->control->step.has_value() ? ctrl->control->step.value() : 1; - auto cloned_value = node_i_let_stmt->value; - auto cloned_let_stmt = LetStmt( - node_i_let_stmt->var.copy_with_suffix(""), - cloned_value, - Evaluate(0) - ); - auto cloned_task = std::make_shared(); - cloned_task->stmts.push_back(cloned_let_stmt); - stage_map[cloned_task.get()] = rem_stage_j; - - for (int k = j; k < n; ++ k) { - auto node_k = seq_body->children[k].get(); - auto task_k = static_cast(node_k); - if (rem_stage_j != stage_map[node_k]) continue; - if (HasDependency(node_i, node_k)) { - for (size_t id = 0; id < task_k->stmts.size(); ++ id) { - task_k->stmts[id] = Substitute( - task_k->stmts[id], - {{node_i_let_stmt->var, cloned_let_stmt->var}} - ); + auto SolveConflictVar = [&]() -> bool { + for (int i = 0; i < n; ++i) + if (IsVarDecl(seq_body->children[i].get())) { + for (int j = 0; j < n; ++j) { + if (i == j) + continue; + + auto node_i = seq_body->children[i].get(); + auto node_j = seq_body->children[j].get(); + int rem_stage_j = stage_map[node_j]; + + if (!HasDependency(node_i, node_j)) + continue; + + // LOG(INFO) << "[ScheduleRecursive] Conflict var detection between + // " << i << " and " << j; + + if (stage_map[node_j] == stage_map[node_i]) + continue; + + auto node_i_task = static_cast(node_i); + auto node_i_let_stmt = node_i_task->stmts[0].as(); + + auto iter = ctrl->control->loop_var; + auto step = ctrl->control->step.has_value() + ? ctrl->control->step.value() + : 1; + auto cloned_value = node_i_let_stmt->value; + auto cloned_let_stmt = + LetStmt(node_i_let_stmt->var.copy_with_suffix(""), cloned_value, + Evaluate(0)); + auto cloned_task = std::make_shared(); + cloned_task->stmts.push_back(cloned_let_stmt); + stage_map[cloned_task.get()] = rem_stage_j; + + for (int k = j; k < n; ++k) { + auto node_k = seq_body->children[k].get(); + auto task_k = static_cast(node_k); + if (rem_stage_j != stage_map[node_k]) + continue; + if (HasDependency(node_i, node_k)) { + for (size_t id = 0; id < task_k->stmts.size(); ++id) { + task_k->stmts[id] = Substitute( + task_k->stmts[id], + {{node_i_let_stmt->var, cloned_let_stmt->var}}); + } + task_k->SubstituteVar(node_i_let_stmt->var, + cloned_let_stmt->var); + stage_map[node_k] = rem_stage_j; } - task_k->SubstituteVar(node_i_let_stmt->var, cloned_let_stmt->var); - stage_map[node_k] = rem_stage_j; } - } - seq_body->children.insert(seq_body->children.begin() + j, std::move(cloned_task)); - n += 1; - return true; // Conflict resolved, restart the loop + seq_body->children.insert(seq_body->children.begin() + j, + std::move(cloned_task)); + n += 1; + return true; // Conflict resolved, restart the loop + } } - } return false; }; - // Resolve conflicts until no more conflicts exist or max iterations reached (to avoid infinite loop) + // Resolve conflicts until no more conflicts exist or max iterations reached + // (to avoid infinite loop) int conflict_count = 0; - while (SolveConflictVar() && ++ conflict_count < 100); + while (SolveConflictVar() && ++conflict_count < 100) + ; for (auto &node : seq_body->children) { auto unit = std::make_shared(); @@ -940,9 +955,7 @@ class ScheduleUnitBuilder { } // Check if two variables are the same - bool SameVar(const Var &a, const Var &b) const { - return a.same_as(b); - } + bool SameVar(const Var &a, const Var &b) const { return a.same_as(b); } // Check if two IRStructures have data dependency (excluding read-after-read) bool HasDependency(const IRStructure *a, const IRStructure *b) const { @@ -1096,8 +1109,9 @@ void ScheduleUnitBuilder::ScheduleRecursive(IRStructure *node) { if (!node) return; - auto ChildrenScheduleHelper = [&](std::vector> origin_children) - -> std::vector> { + auto ChildrenScheduleHelper = + [&](std::vector> origin_children) + -> std::vector> { // Now collect child nodes for potential scheduling std::vector child_nodes; child_nodes.reserve(origin_children.size()); @@ -1181,7 +1195,7 @@ void ScheduleUnitBuilder::ScheduleRecursive(IRStructure *node) { Z3SchedulePythonLoop(ctrl); } else if (ctrl->child->IsWrapper()) { auto wrapper = static_cast(ctrl->child.get()); - std::vector>origin_children; + std::vector> origin_children; GatherTaskNodes({wrapper->child}, origin_children); for (auto &child : origin_children) { ScheduleRecursive(child.get()); @@ -1194,7 +1208,7 @@ void ScheduleUnitBuilder::ScheduleRecursive(IRStructure *node) { return; } else if (node->IsWrapper()) { auto wrapper = static_cast(node); - std::vector>origin_children; + std::vector> origin_children; GatherTaskNodes({wrapper->child}, origin_children); for (auto &child : origin_children) { ScheduleRecursive(child.get()); @@ -1434,7 +1448,8 @@ class IRStructureBuilder : public StmtVisitor { int64_t thread_count_ = 1; Target target_; - void AnalyzeResourceUsage(const Stmt &stmt, TaskNode *task_node, bool only_variables = false) { + void AnalyzeResourceUsage(const Stmt &stmt, TaskNode *task_node, + bool only_variables = false) { // Recursively analyze statements to determine resource usage struct ResourceAnalyzer : public StmtExprVisitor { TaskNode *task_node; @@ -1606,8 +1621,8 @@ class IRStructureBuilder : public StmtVisitor { task_node->SetGemmInst(analyzer.gemm_inst); } } - // If neither TMA nor Tensor core was used, and CUDA operations were found, - // set CUDA core flag + // If neither TMA nor Tensor core was used, and CUDA operations were + // found, set CUDA core flag if (!analyzer.found_tma && !analyzer.found_tensor) { task_node->SetUsesCUDACore(true); } @@ -1887,14 +1902,14 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) { // Helper: check if a TaskNode is a LetDecl (single LetStmt with empty body) static bool IsLetDeclTask(const TaskNode *task) { - return task->stmts.size() == 1 && - task->stmts[0].as() != nullptr; + return task->stmts.size() == 1 && task->stmts[0].as() != nullptr; } // Helper: check if an IRStructure node is a LetDecl task (or a ScheduleUnit // wrapping one) static bool IsLetDeclNode(const IRStructure *node) { - if (!node) return false; + if (!node) + return false; if (node->IsTask()) { return IsLetDeclTask(static_cast(node)); } @@ -1908,12 +1923,15 @@ static bool IsLetDeclNode(const IRStructure *node) { // Helper: check if an IRStructure subtree contains any LetDecl tasks static bool ContainsLetDecl(const IRStructure *node) { - if (!node) return false; - if (IsLetDeclNode(node)) return true; + if (!node) + return false; + if (IsLetDeclNode(node)) + return true; if (node->IsSequence()) { auto seq = static_cast(node); for (const auto &child : seq->children) { - if (ContainsLetDecl(child.get())) return true; + if (ContainsLetDecl(child.get())) + return true; } } else if (node->IsControl()) { auto ctrl = static_cast(node); @@ -1932,7 +1950,8 @@ static bool ContainsLetDecl(const IRStructure *node) { std::shared_ptr CloneIRStructureWithWarpgroupFilter(IRStructure *node, int warpgroup_id, Map &var_remap) { - if (!node) return nullptr; + if (!node) + return nullptr; if (node->IsTask()) { auto task = static_cast(node); @@ -1943,9 +1962,8 @@ CloneIRStructureWithWarpgroupFilter(IRStructure *node, int warpgroup_id, const auto *let = task->stmts[0].as(); auto new_var = let->var.copy_with_suffix(""); // Substitute previously renamed variables in the value expression. - PrimExpr new_value = var_remap.empty() - ? let->value - : Substitute(let->value, var_remap); + PrimExpr new_value = + var_remap.empty() ? let->value : Substitute(let->value, var_remap); var_remap.Set(let->var, new_var); auto new_task = std::make_shared(); new_task->stmts.push_back(LetStmt(new_var, new_value, Evaluate(0))); @@ -1953,7 +1971,8 @@ CloneIRStructureWithWarpgroupFilter(IRStructure *node, int warpgroup_id, } // Non-LetDecl tasks: only include if warp group matches - if (!node->containWarpgroupId(warpgroup_id)) return nullptr; + if (!node->containWarpgroupId(warpgroup_id)) + return nullptr; auto cloned = task->Clone(); // Substitute renamed LetDecl variables in task statements if (!var_remap.empty()) { @@ -1977,7 +1996,8 @@ CloneIRStructureWithWarpgroupFilter(IRStructure *node, int warpgroup_id, new_seq->children.push_back(std::move(new_child)); } } - if (new_seq->children.empty()) return nullptr; + if (new_seq->children.empty()) + return nullptr; return new_seq; } else if (node->IsControl()) { // A ControlNode is included if it contains the target warp group @@ -2055,7 +2075,8 @@ class VarRefCollector : public StmtExprVisitor { // tasks that used them ended up in the other warp group. std::shared_ptr RemoveUnusedLetDecls(std::shared_ptr root) { - if (!root) return nullptr; + if (!root) + return nullptr; // Phase 1: Collect LetDecl definitions and variable references from // non-LetDecl nodes (task stmts and ScheduleUnit before/after). @@ -2068,7 +2089,8 @@ RemoveUnusedLetDecls(std::shared_ptr root) { std::function collect = [&](const IRStructure *node) { - if (!node) return; + if (!node) + return; if (node->IsTask()) { auto task = static_cast(node); if (IsLetDeclTask(task)) { @@ -2096,13 +2118,14 @@ RemoveUnusedLetDecls(std::shared_ptr root) { collect(unit->child.get()); VarRefCollector collector; for (const auto &stmts : unit->before) { - for (const auto &s : stmts) collector(s); + for (const auto &s : stmts) + collector(s); } for (const auto &stmts : unit->after) { - for (const auto &s : stmts) collector(s); + for (const auto &s : stmts) + collector(s); } - referenced_vars.insert(collector.vars.begin(), - collector.vars.end()); + referenced_vars.insert(collector.vars.begin(), collector.vars.end()); } }; collect(root.get()); @@ -2131,13 +2154,13 @@ RemoveUnusedLetDecls(std::shared_ptr root) { const std::shared_ptr &)> filter_tree = [&](const std::shared_ptr &node) -> std::shared_ptr { - if (!node) return nullptr; + if (!node) + return nullptr; if (node->IsTask()) { if (IsLetDeclTask(static_cast(node.get()))) { - const auto *let = - static_cast(node.get()) - ->stmts[0] - .as(); + const auto *let = static_cast(node.get()) + ->stmts[0] + .as(); if (!referenced_vars.count(let->var.get())) { return nullptr; // Remove unused LetDecl } @@ -2148,9 +2171,11 @@ RemoveUnusedLetDecls(std::shared_ptr root) { auto new_seq = std::make_shared(); for (const auto &child : seq->children) { auto filtered = filter_tree(child); - if (filtered) new_seq->children.push_back(std::move(filtered)); + if (filtered) + new_seq->children.push_back(std::move(filtered)); } - if (new_seq->children.empty()) return nullptr; + if (new_seq->children.empty()) + return nullptr; return new_seq; } else if (node->IsControl()) { auto ctrl = static_cast(node.get()); @@ -2158,7 +2183,8 @@ RemoveUnusedLetDecls(std::shared_ptr root) { new_ctrl->control = ctrl->control; new_ctrl->SetPromote(ctrl->hasPromote()); new_ctrl->child = filter_tree(ctrl->child); - if (!new_ctrl->child) return nullptr; + if (!new_ctrl->child) + return nullptr; return new_ctrl; } else if (node->IsWrapper()) { auto wrapper = static_cast(node.get()); @@ -2173,7 +2199,8 @@ RemoveUnusedLetDecls(std::shared_ptr root) { new_unit->before = unit->before; new_unit->after = unit->after; new_unit->child = filter_tree(unit->child); - if (!new_unit->child) return nullptr; + if (!new_unit->child) + return nullptr; return new_unit; } return node; @@ -2387,8 +2414,11 @@ Stmt ConvertIRStructureToStmt(IRStructure *root, const bool outer_enable_epi) { Map substitution, substitution_cond; substitution.Set(loop_var, loop_var - loop_step * (max_stages - unit->stage)); - substitution_cond.Set(loop_var, - Max(loop_start, Min(loop_start + loop_extent - loop_step, loop_var - loop_step * (max_stages - unit->stage)))); + substitution_cond.Set( + loop_var, + Max(loop_start, + Min(loop_start + loop_extent - loop_step, + loop_var - loop_step * (max_stages - unit->stage)))); if (IsLetDeclNode(unit->child.get())) { Stmt stmt = SeqStmt::Flatten(stmts); steady.push_back(Substitute(stmt, substitution_cond)); @@ -2702,8 +2732,11 @@ Stmt ApplyWarpgroupPartitionToIRStructure( Map substitution, substitution_cond; substitution.Set(loop_var, loop_var - loop_step * (max_stages - unit->stage)); - substitution_cond.Set(loop_var, - Max(loop_start, Min(loop_start + loop_extent - loop_step, loop_var - loop_step * (max_stages - unit->stage)))); + substitution_cond.Set( + loop_var, + Max(loop_start, + Min(loop_start + loop_extent - loop_step, + loop_var - loop_step * (max_stages - unit->stage)))); if (IsLetDeclNode(unit->child.get())) { Stmt stmt = SeqStmt::Flatten(stmts); steady.push_back(Substitute(stmt, substitution_cond)); @@ -3111,23 +3144,27 @@ class LetStmtNester : public StmtMutator { } stmts = flat_stmts; - for (int i = static_cast(stmts.size()) - 2; i >= 0; -- i) { + for (int i = static_cast(stmts.size()) - 2; i >= 0; --i) { if (const auto *let = stmts[i].as()) { if (IsEmptyBody(let->body)) { Stmt absorbed_body = CollectRemaining(stmts, i + 1); - stmts = TruncateAndReplace(stmts, i, LetStmt(let->var, let->value, absorbed_body)); + stmts = TruncateAndReplace( + stmts, i, LetStmt(let->var, let->value, absorbed_body)); } } else if (const auto *attr = stmts[i].as()) { if (IsEmptyBody(attr->body)) { Stmt absorbed_body = CollectRemaining(stmts, i + 1); - stmts = TruncateAndReplace(stmts, i, - AttrStmt(attr->node, attr->attr_key, attr->value, absorbed_body)); + stmts = TruncateAndReplace( + stmts, i, + AttrStmt(attr->node, attr->attr_key, attr->value, absorbed_body)); } } } - if (stmts.empty()) return Evaluate(0); - if (stmts.size() == 1) return stmts[0]; + if (stmts.empty()) + return Evaluate(0); + if (stmts.size() == 1) + return stmts[0]; return SeqStmt(stmts); } diff --git a/src/transform/auto_schedule/ir_structure.cc b/src/transform/auto_schedule/ir_structure.cc index b23330e126..394e776307 100644 --- a/src/transform/auto_schedule/ir_structure.cc +++ b/src/transform/auto_schedule/ir_structure.cc @@ -122,7 +122,8 @@ std::vector SequenceNode::GetReadVars() const { for (const auto &child : children) { if (child) { auto child_read_vars = child->GetReadVars(); - all_vars.insert(all_vars.end(), child_read_vars.begin(), child_read_vars.end()); + all_vars.insert(all_vars.end(), child_read_vars.begin(), + child_read_vars.end()); } } return all_vars; @@ -160,7 +161,8 @@ std::vector SequenceNode::GetWriteVars() const { for (const auto &child : children) { if (child) { auto child_write_vars = child->GetWriteVars(); - all_vars.insert(all_vars.end(), child_write_vars.begin(), child_write_vars.end()); + all_vars.insert(all_vars.end(), child_write_vars.begin(), + child_write_vars.end()); } } return all_vars; diff --git a/src/transform/auto_schedule/ir_structure.h b/src/transform/auto_schedule/ir_structure.h index 28ba5711bf..6df85be596 100644 --- a/src/transform/auto_schedule/ir_structure.h +++ b/src/transform/auto_schedule/ir_structure.h @@ -101,7 +101,6 @@ class IRStructure { virtual void SetLatency(int64_t latency) = 0; virtual void SetII(int64_t ii) = 0; - // Recursive region collection method virtual void CollectRegions( std::vector &result, @@ -144,16 +143,12 @@ class TaskNode : public IRStructure { std::vector GetWriteRegions() const override { return write_regions_; } - std::vector GetReadVars() const override { - return read_vars_; - } - std::vector GetWriteVars() const override { - return write_vars_; - } + std::vector GetReadVars() const override { return read_vars_; } + std::vector GetWriteVars() const override { return write_vars_; } // Latency estimation int64_t GetLatency() const override { return latency_; } int64_t GetII() const override { return ii_; } - + // Setters void SetUsesCUDACore(bool value) override { uses_cuda_core_ = value; } void SetUsesTMACore(bool value) override { uses_tma_core_ = value; } @@ -164,12 +159,8 @@ class TaskNode : public IRStructure { void SetWriteRegions(const std::vector ®ions) override { write_regions_ = regions; } - void SetReadVars(const std::vector &vars) { - read_vars_ = vars; - } - void SetWriteVars(const std::vector &vars) { - write_vars_ = vars; - } + void SetReadVars(const std::vector &vars) { read_vars_ = vars; } + void SetWriteVars(const std::vector &vars) { write_vars_ = vars; } void SubstituteVar(const Var &old_var, const Var &new_var) { for (auto &var : read_vars_) { if (var.same_as(old_var)) { @@ -303,7 +294,9 @@ class TaskNode : public IRStructure { std::vector &result, std::set>> &visited) const override; - bool containWarpgroupId(int id) const override { return ContainsLoopBreak() || warpgroup_id_ == id; } + bool containWarpgroupId(int id) const override { + return ContainsLoopBreak() || warpgroup_id_ == id; + } // Check if this task contains loop_break call bool ContainsLoopBreak() const; @@ -876,7 +869,8 @@ inline Stmt GetAttrDecl(const AttrStmtNode *attr_node) { if (attr_node == nullptr) { return Stmt(); } - return AttrStmt(attr_node->node, attr_node->attr_key, attr_node->value, Evaluate(0)); + return AttrStmt(attr_node->node, attr_node->attr_key, attr_node->value, + Evaluate(0)); } // Original helper function to print IRStructure (kept for backward diff --git a/src/transform/auto_schedule/memory_detector.h b/src/transform/auto_schedule/memory_detector.h index 0364b2af20..8a37fc6eea 100644 --- a/src/transform/auto_schedule/memory_detector.h +++ b/src/transform/auto_schedule/memory_detector.h @@ -126,7 +126,8 @@ class MemoryAccessDetector : public StmtExprVisitor { */ void UpdateReadVar(const Var &var) { for (const auto &v : read_vars_) { - if (v.same_as(var)) return; + if (v.same_as(var)) + return; } read_vars_.push_back(var); } @@ -137,7 +138,8 @@ class MemoryAccessDetector : public StmtExprVisitor { */ void UpdateWriteVar(const Var &var) { for (const auto &v : write_vars_) { - if (v.same_as(var)) return; + if (v.same_as(var)) + return; } write_vars_.push_back(var); } diff --git a/src/transform/if_condition_extract.cc b/src/transform/if_condition_extract.cc index 9df3de18f5..bca573b348 100644 --- a/src/transform/if_condition_extract.cc +++ b/src/transform/if_condition_extract.cc @@ -1,6 +1,7 @@ /*! * \file if_condition_extract.cc - * \brief Extract if conditions into temporary LetStmt variables, then expand if statements to all branches. + * \brief Extract if conditions into temporary LetStmt variables, then expand if + * statements to all branches. */ #include @@ -58,13 +59,15 @@ class IfConditionExtractor : public StmtExprMutator { auto bind_cond_var = [](const Stmt &sentence, const Var &cond) -> Stmt { if (auto if_sentence = sentence.as()) { PrimExpr new_cond = cond & if_sentence->condition; - return IfThenElse(new_cond, if_sentence->then_case, if_sentence->else_case); + return IfThenElse(new_cond, if_sentence->then_case, + if_sentence->else_case); } else { return IfThenElse(cond, sentence); } }; - auto bind_cond_var_body = [&](const Optional &body, const Var &cond) -> Stmt { + auto bind_cond_var_body = [&](const Optional &body, + const Var &cond) -> Stmt { if (!body.defined()) { return Stmt(); } @@ -81,9 +84,13 @@ class IfConditionExtractor : public StmtExprMutator { Array new_seq; new_seq.insert(new_seq.end(), bind_cond_var_body(then_case, cond_var)); - if (else_case.defined()) new_seq.insert(new_seq.end(), bind_cond_var_body(else_case, cond_var)); + if (else_case.defined()) + new_seq.insert(new_seq.end(), bind_cond_var_body(else_case, cond_var)); - Stmt body = new_seq.empty() ? Stmt() : (new_seq.size() == 1 ? new_seq[0] : SeqStmt(std::move(new_seq))); + Stmt body = + new_seq.empty() + ? Stmt() + : (new_seq.size() == 1 ? new_seq[0] : SeqStmt(std::move(new_seq))); if (is_simple) { return body; } else { @@ -95,7 +102,8 @@ class IfConditionExtractor : public StmtExprMutator { Array seq; for (auto stmt : op->seq) { auto new_stmt = VisitStmt(stmt); - if (!new_stmt.defined()) continue; + if (!new_stmt.defined()) + continue; if (auto seq_node = new_stmt.as()) { seq.insert(seq.end(), seq_node->seq.begin(), seq_node->seq.end()); } else { diff --git a/tilelang/transform/__init__.py b/tilelang/transform/__init__.py index a6e97dd023..689258cca0 100644 --- a/tilelang/transform/__init__.py +++ b/tilelang/transform/__init__.py @@ -561,10 +561,12 @@ def LowerSharedTmem(): """LowerSharedTmem""" return _ffi_api.LowerSharedTmem() # type: ignore + def IfConditionExtract(): """Extract if condition to LetStmt variables.""" return _ffi_api.IfConditionExtract() # type: ignore + def AutoSchedule(enable_epi: bool): """Auto schedule for high-level operations""" return _ffi_api.AutoSchedule(enable_epi) # type: ignore