diff --git a/benchmark/gemm/frost/benchmark_block_scale_matmul.py b/benchmark/gemm/frost/benchmark_block_scale_matmul.py index 7025e9280e..d277f919ab 100644 --- a/benchmark/gemm/frost/benchmark_block_scale_matmul.py +++ b/benchmark/gemm/frost/benchmark_block_scale_matmul.py @@ -145,8 +145,11 @@ def _mkdata(batch: int, M: int, N: int, K: int, combo: str): sfb_log = _rand_e8m0((batch, N, sf_k), dev) c = torch.empty(batch, M, N, dtype=torch.bfloat16, device=dev) - sfa = torch.cat([_to_blocked(sfa_log[i]) for i in range(batch)]).view(batch, M, sf_k) - sfb = torch.cat([_to_blocked(sfb_log[i]) for i in range(batch)]).view(batch, N, sf_k) + # The blocked blob carries the PADDED dims (whole 128-row x 4-SF-K atoms) — + # that is what the kernel reads. The graph keeps the logical [batch, M, sf_k]. + sf_k_pad = _ceil_div(sf_k, 4) * 4 + sfa = torch.cat([_to_blocked(sfa_log[i]) for i in range(batch)]).view(batch, _ceil_div(M, 128) * 128, sf_k_pad) + sfb = torch.cat([_to_blocked(sfb_log[i]) for i in range(batch)]).view(batch, _ceil_div(N, 128) * 128, sf_k_pad) return a, b, c, sfa, sfb @@ -200,7 +203,7 @@ def _rotating(fn_of_set: Callable, pool: list) -> Callable: return lambda i: fn_of_set(pool[i % n]) -def _scaled_mm_ref(batch: int, M: int, N: int, K: int, combo: str): +def _scaled_mm_ref(batch: int, M: int, N: int, K: int, combo: str, verbose: bool = True): """(label, call) for cuBLAS's OWN block-scaled GEMM of this combo, via torch.nn.functional.scaled_mm (cuBLASLt) — the apples-to-apples reference. Returns None if the env can't run it (some cuBLASLt builds return no @@ -217,7 +220,6 @@ def _scaled_mm_ref(batch: int, M: int, N: int, K: int, combo: str): is_fp4, bs, _, _ = _COMBOS[combo] ru = lambda x, m: ((x + m - 1) // m) * m cd = lambda a, b: (a + b - 1) // b - sw = [SwizzleType.SWIZZLE_32_4_4] if is_fp4: a = torch.randint(0, 256, (M, K // 2), dtype=torch.uint8, device=dev).view(torch.float4_e2m1fn_x2) @@ -235,11 +237,14 @@ def _scaled_mm_ref(batch: int, M: int, N: int, K: int, combo: str): gb = torch.ones(1, device=dev) recipe = [ScalingType.BlockWise1x16, ScalingType.TensorWise] scA, scB = [sa, ga], [sb, gb] + # scaled_mm wants one swizzle PER scale in the recipe. + sw = [SwizzleType.SWIZZLE_32_4_4, SwizzleType.NO_SWIZZLE] else: # mxfp4 / mxfp8 — no global scale sa = torch.randn(na, device=dev).to(torch.float8_e8m0fnu) sb = torch.randn(nb, device=dev).to(torch.float8_e8m0fnu) recipe = [ScalingType.BlockWise1x32] scA, scB = [sa], [sb] + sw = [SwizzleType.SWIZZLE_32_4_4] # scaled_mm has no batched form, so a batched ref is B back-to-back single-GEMM # calls (same operands); FLOPS counts all B so throughput stays comparable. @@ -252,7 +257,10 @@ def call(): try: call() torch.cuda.synchronize() - except Exception: + except Exception as e: + # Print it — an unrunnable env and an API-shape bug look identical otherwise. + if verbose: + print(f" [scaled_mm '{combo}' reference unavailable: {type(e).__name__}: {e}]") return None suffix = f" ×{batch}" if batch > 1 else "" return f"cuBLAS {combo} scaled_mm{suffix}", call @@ -265,15 +273,15 @@ def _make_reference(batch: int, M: int, N: int, K: int, combo: str, ref_mode: st / bf16 (dense BF16) / auto (scaled_mm, falling back to BF16).""" dev = "cuda" if ref_mode in ("auto", "scaled_mm"): - ref = _scaled_mm_ref(batch, M, N, K, combo) + ref = _scaled_mm_ref(batch, M, N, K, combo, verbose=verbose) if ref is not None: return ref if ref_mode == "scaled_mm": sys.exit( f"--ref scaled_mm: cuBLAS block-scaled GEMM for '{combo}' is not " - f"runnable here — torch.nn.functional.scaled_mm raised at the " - f"cuBLASLt heuristic (no algorithm for this driver / cuBLASLt " - f"build). Use --ref bf16, or --ref auto to fall back automatically." + f"runnable here — torch.nn.functional.scaled_mm raised (see the " + f"message above). Use --ref bf16, or --ref auto to fall back " + f"automatically." ) if verbose: print(" [note: scaled_mm reference unavailable in this env — " "falling back to dense BF16 cuBLAS]") diff --git a/benchmark/gemm/frost/benchmark_block_scale_matmul_swiglu.py b/benchmark/gemm/frost/benchmark_block_scale_matmul_swiglu.py index b45784fe4b..f90f392635 100644 --- a/benchmark/gemm/frost/benchmark_block_scale_matmul_swiglu.py +++ b/benchmark/gemm/frost/benchmark_block_scale_matmul_swiglu.py @@ -141,7 +141,8 @@ def _mk(rows): def _sf(rows): log = torch.randint(1, 4, (B, rows, sf_k), device=dev).to(torch.float8_e4m3fn) - blk = torch.stack([_to_blocked(log[b]) for b in range(B)]).view(B, rows, sf_k) + # _to_blocked pads to whole 128-row x 4-SF-K atoms — view the PADDED dims. + blk = torch.stack([_to_blocked(log[b]) for b in range(B)]).view(B, _ceil_div(rows, 128) * 128, _ceil_div(sf_k, 4) * 4) return log, blk sfa_log, sfa_b = _sf(M) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma.py index d4dae8079f..65e26d2e5b 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma.py @@ -116,6 +116,7 @@ def _kernel( warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) + elect_one = nvvm.elect_sync() tidx = cute.arch.thread_idx()[0] bidx = cute.arch.block_idx()[0] @@ -252,7 +253,7 @@ def _kernel( num_consumer_warps_per_cta = 7 clc_empty_count = num_consumer_warps_per_cta * cluster_size if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: for i in range(ab_stages): nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) nvvm.mbarrier_init(ab_empty_mbar_ptr.subview(i), ab_empty_count) @@ -305,11 +306,11 @@ def _kernel( while not nvvm.mbarrier_try_wait_parity(clc_empty_mbar_ptr.subview(stage), clc_empty_phase, time_limit=10_000_000): pass - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(clc_full_mbar_ptr.subview(stage), 16) if is_cluster_leader_cta: - if nvvm.elect_sync(): + if elect_one: cute_clc.issue_clc_query( clc_full_mbar_cute_base + stage, clc_response_ptr_base + stage, @@ -324,7 +325,7 @@ def _kernel( is_valid_sched = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) @@ -347,7 +348,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) @@ -379,7 +380,7 @@ def _kernel( coord_k = k_tile_idx * cta_tile_mnk[2] coord_sf_k = k_tile_idx * sf_tma_box_k - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) for _ai in cutlass.range_constexpr(num_a_operands): @@ -390,7 +391,7 @@ def _kernel( sfa_m_block = coord_m_per_cta // 128 if cutlass.const_expr(multicast_a): if n_rank == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -417,7 +418,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_1, ) else: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -445,7 +446,7 @@ def _kernel( ) if cutlass.const_expr(multicast_a): if n_rank == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFA_stage, tma_sfa_desc.get_ptr(), @@ -456,7 +457,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_1, ) else: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFA_stage, tma_sfa_desc.get_ptr(), @@ -475,7 +476,7 @@ def _kernel( sfb_n_block = coord_n_per_cta // 128 if cutlass.const_expr(multicast_b): if m_rank == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -502,7 +503,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_1, ) else: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -530,7 +531,7 @@ def _kernel( ) if cutlass.const_expr(multicast_b): if m_rank == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFB_stage, tma_sfb_desc.get_ptr(), @@ -541,7 +542,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_1, ) else: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFB_stage, tma_sfb_desc.get_ptr(), @@ -575,7 +576,7 @@ def _kernel( ) tile_l = l_idx nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 @@ -589,7 +590,7 @@ def _kernel( if tail_stage == ab_stages: tail_stage = cutlass.Int32(0) tail_phase = tail_phase ^ 1 - if nvvm.elect_sync(): + if elect_one: while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): pass @@ -706,7 +707,7 @@ def _kernel( for atom_r in cutlass.range(num_sf_atoms, unroll_full=True): for _ai in cutlass.range_constexpr(num_a_operands): for _m in cutlass.range_constexpr(num_blocks_m): - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_cp( s2t_shape, sfa_dst_ptrs[_ai][_m], @@ -716,7 +717,7 @@ def _kernel( ) for _bj in cutlass.range_constexpr(num_b_operands): for _m in cutlass.range_constexpr(num_blocks_n): - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_cp( s2t_shape, sfb_dst_ptrs[_bj][_m], @@ -732,7 +733,7 @@ def _kernel( _bj = gemm_b_idx[g] desc_a = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * k_block_idx) desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * k_block_idx) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_1, @@ -747,7 +748,7 @@ def _kernel( ) scale_d = cutlass.Boolean(True) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( ab_empty_mbar_ptr.subview(stage), multicast_mask=ab_empty_arrive_mask, @@ -755,7 +756,7 @@ def _kernel( ) ab_iter += 1 - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( acc_full_mbar_ptr.subview(acc_stage), group=nvvm.CTAGroup.CTA_1, @@ -774,18 +775,18 @@ def _kernel( cute.arch.fence_proxy("async.shared", space="cta") is_valid = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if nvvm.elect_sync(): + if elect_one: for _ in range(acc_stages): tail_stage = tail_stage + 1 if tail_stage == acc_stages: @@ -878,7 +879,7 @@ def _kernel( if use_acc_overlap and (not cd_out_is_m_major) and subtile_idx == acc_overlap_subtiles - 1: nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) col = coord_n + subtile_col_offset @@ -934,7 +935,7 @@ def _kernel( ) if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(cd_out_is_m_major): for _mb in cutlass.range_constexpr(cta_tile_mnk[0] // cd_mmajor_atom_m): nvvm.cp_async_bulk_tensor_global_shared_cta( @@ -971,7 +972,7 @@ def _kernel( if cutlass.const_expr((not use_acc_overlap) or cd_out_is_m_major): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) consumer_stage = tile_iter % CLC_SCHED_STAGES @@ -995,7 +996,7 @@ def _kernel( ) tile_l = l_idx nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) @@ -1003,7 +1004,7 @@ def _kernel( if cutlass.const_expr(use_acc_overlap): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(tmem_dealloc_mbar_ptr) # @@TMA_STORE_ONLY:BEGIN@@ diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma_static.py b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma_static.py index 9998f5cd36..a7e0bda1d8 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma_static.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma_static.py @@ -115,6 +115,7 @@ def _kernel( warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) + elect_one = nvvm.elect_sync() tidx = cute.arch.thread_idx()[0] bidx = cute.arch.block_idx()[0] @@ -236,7 +237,7 @@ def _kernel( ab_empty_count = cluster_m + cluster_n - 1 if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: for i in range(ab_stages): nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) nvvm.mbarrier_init(ab_empty_mbar_ptr.subview(i), ab_empty_count) @@ -278,7 +279,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) @@ -309,7 +310,7 @@ def _kernel( coord_k = k_tile_idx * cta_tile_mnk[2] coord_sf_k = k_tile_idx * sf_tma_box_k - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) for _ai in cutlass.range_constexpr(num_a_operands): @@ -320,7 +321,7 @@ def _kernel( sfa_m_block = coord_m_per_cta // 128 if cutlass.const_expr(multicast_a): if n_rank == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -347,7 +348,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_1, ) else: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -375,7 +376,7 @@ def _kernel( ) if cutlass.const_expr(multicast_a): if n_rank == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFA_stage, tma_sfa_desc.get_ptr(), @@ -386,7 +387,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_1, ) else: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFA_stage, tma_sfa_desc.get_ptr(), @@ -405,7 +406,7 @@ def _kernel( sfb_n_block = coord_n_per_cta // 128 if cutlass.const_expr(multicast_b): if m_rank == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -432,7 +433,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_1, ) else: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -460,7 +461,7 @@ def _kernel( ) if cutlass.const_expr(multicast_b): if m_rank == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFB_stage, tma_sfb_desc.get_ptr(), @@ -471,7 +472,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_1, ) else: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFB_stage, tma_sfb_desc.get_ptr(), @@ -496,7 +497,7 @@ def _kernel( if tail_stage == ab_stages: tail_stage = cutlass.Int32(0) tail_phase = tail_phase ^ 1 - if nvvm.elect_sync(): + if elect_one: while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): pass @@ -614,7 +615,7 @@ def _kernel( for atom_r in cutlass.range(num_sf_atoms, unroll_full=True): for _ai in cutlass.range_constexpr(num_a_operands): for _m in cutlass.range_constexpr(num_blocks_m): - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_cp( s2t_shape, sfa_dst_ptrs[_ai][_m], @@ -624,7 +625,7 @@ def _kernel( ) for _bj in cutlass.range_constexpr(num_b_operands): for _m in cutlass.range_constexpr(num_blocks_n): - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_cp( s2t_shape, sfb_dst_ptrs[_bj][_m], @@ -640,7 +641,7 @@ def _kernel( _bj = gemm_b_idx[g] desc_a = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * k_block_idx) desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * k_block_idx) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_1, @@ -655,7 +656,7 @@ def _kernel( ) scale_d = cutlass.Boolean(True) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( ab_empty_mbar_ptr.subview(stage), multicast_mask=ab_empty_arrive_mask, @@ -663,7 +664,7 @@ def _kernel( ) ab_iter += 1 - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( acc_full_mbar_ptr.subview(acc_stage), group=nvvm.CTAGroup.CTA_1, @@ -673,12 +674,12 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) - if nvvm.elect_sync(): + if elect_one: while not nvvm.mbarrier_try_wait_parity( acc_empty_mbar_ptr.subview(acc_stage), acc_empty_phase_bit ^ 1, @@ -772,7 +773,7 @@ def _kernel( if use_acc_overlap and (not cd_out_is_m_major) and subtile_idx == acc_overlap_subtiles - 1: nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) col = coord_n + subtile_col_offset @@ -828,7 +829,7 @@ def _kernel( ) if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(cd_out_is_m_major): for _mb in cutlass.range_constexpr(cta_tile_mnk[0] // cd_mmajor_atom_m): nvvm.cp_async_bulk_tensor_global_shared_cta( @@ -865,7 +866,7 @@ def _kernel( if cutlass.const_expr((not use_acc_overlap) or cd_out_is_m_major): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) is_valid = cutlass.Int32(0) @@ -874,7 +875,7 @@ def _kernel( if cutlass.const_expr(use_acc_overlap): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(tmem_dealloc_mbar_ptr) # @@TMA_STORE_ONLY:BEGIN@@ diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma.py index 12b13e7e34..4bee489e3f 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma.py @@ -116,6 +116,7 @@ def _kernel( warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) + elect_one = nvvm.elect_sync() tidx = cute.arch.thread_idx()[0] bidx = cute.arch.block_idx()[0] @@ -254,7 +255,7 @@ def _kernel( num_consumer_warps_per_cta = 7 clc_empty_count = num_consumer_warps_per_cta * cluster_size if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(use_acc_overlap): nvvm.mbarrier_init(tmem_dealloc_mbar_ptr, num_epilogue_warps) else: @@ -312,11 +313,11 @@ def _kernel( while not nvvm.mbarrier_try_wait_parity(clc_empty_mbar_ptr.subview(stage), clc_empty_phase, time_limit=10_000_000): pass - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(clc_full_mbar_ptr.subview(stage), 16) if is_cluster_leader_cta: - if nvvm.elect_sync(): + if elect_one: cute_clc.issue_clc_query( clc_full_mbar_cute_base + stage, clc_response_ptr_base + stage, @@ -331,7 +332,7 @@ def _kernel( is_valid_sched = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) @@ -354,7 +355,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) @@ -391,7 +392,7 @@ def _kernel( sfb_n_block = coord_n_pair // 128 if is_pair_leader: - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) for _ai in cutlass.range_constexpr(num_a_operands): @@ -402,7 +403,7 @@ def _kernel( sfa_m_block = coord_m_per_cta // 128 if cutlass.const_expr(multicast_a): if n_rank == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -429,7 +430,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_2, ) else: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -457,7 +458,7 @@ def _kernel( ) if cutlass.const_expr(multicast_a): if n_rank == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFA_stage, tma_sfa_desc.get_ptr(), @@ -468,7 +469,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_2, ) else: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFA_stage, tma_sfa_desc.get_ptr(), @@ -486,7 +487,7 @@ def _kernel( tma_sfb_desc = tma_sfb_descs[_bj] if cutlass.const_expr(multicast_b): if pair_m_idx == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -513,7 +514,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_2, ) else: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -541,7 +542,7 @@ def _kernel( ) if cutlass.const_expr(multicast_b): if pair_m_idx == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFB_stage, tma_sfb_desc.get_ptr(), @@ -552,7 +553,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_2, ) else: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFB_stage, tma_sfb_desc.get_ptr(), @@ -586,7 +587,7 @@ def _kernel( ) tile_l = l_idx nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 @@ -600,7 +601,7 @@ def _kernel( if tail_stage == ab_stages: tail_stage = cutlass.Int32(0) tail_phase = tail_phase ^ 1 - if nvvm.elect_sync(): + if elect_one: while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): pass @@ -736,7 +737,7 @@ def _kernel( for atom_r in cutlass.range(num_sf_atoms, unroll_full=True): for _ai in cutlass.range_constexpr(num_a_operands): for _m in cutlass.range_constexpr(num_blocks_m): - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_cp( s2t_shape, sfa_dst_ptrs[_ai][_m], @@ -746,7 +747,7 @@ def _kernel( ) for _bj in cutlass.range_constexpr(num_b_operands): for _m in cutlass.range_constexpr(num_blocks_n): - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_cp( s2t_shape, sfb_dst_ptrs[_bj][_m], @@ -762,7 +763,7 @@ def _kernel( _bj = gemm_b_idx[g] desc_a = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * k_block_idx) desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * k_block_idx) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_2, @@ -777,7 +778,7 @@ def _kernel( ) scale_d = cutlass.Boolean(True) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( ab_empty_mbar_ptr.subview(stage), multicast_mask=ab_empty_arrive_mask, @@ -785,7 +786,7 @@ def _kernel( ) ab_iter += 1 - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( acc_full_mbar_ptr.subview(acc_stage), multicast_mask=pair_mask, @@ -805,18 +806,18 @@ def _kernel( cute.arch.fence_proxy("async.shared", space="cta") is_valid = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if nvvm.elect_sync(): + if elect_one: for _ in range(acc_stages): tail_stage = tail_stage + 1 if tail_stage == acc_stages: @@ -856,13 +857,13 @@ def _kernel( cute.arch.fence_proxy("async.shared", space="cta") is_valid = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) @@ -956,7 +957,7 @@ def _kernel( if use_acc_overlap and (not cd_out_is_m_major) and subtile_idx == acc_overlap_subtiles - 1: nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) @@ -1013,7 +1014,7 @@ def _kernel( ) if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(cd_out_is_m_major): for _mb in cutlass.range_constexpr(cta_tile_mnk[0] // cd_mmajor_atom_m): nvvm.cp_async_bulk_tensor_global_shared_cta( @@ -1050,7 +1051,7 @@ def _kernel( if cutlass.const_expr((not use_acc_overlap) or cd_out_is_m_major): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) @@ -1075,7 +1076,7 @@ def _kernel( ) tile_l = l_idx nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) @@ -1083,7 +1084,7 @@ def _kernel( if cutlass.const_expr(use_acc_overlap): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(tmem_dealloc_mbar_ptr) # @@TMA_STORE_ONLY:BEGIN@@ diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma_static.py b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma_static.py index ee1a4ec101..7c5a154977 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma_static.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma_static.py @@ -113,6 +113,7 @@ def _kernel( warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) + elect_one = nvvm.elect_sync() tidx = cute.arch.thread_idx()[0] bidx = cute.arch.block_idx()[0] @@ -235,7 +236,7 @@ def _kernel( cta_group = 2 ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(use_acc_overlap): nvvm.mbarrier_init(tmem_dealloc_mbar_ptr, num_epilogue_warps) else: @@ -280,7 +281,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) @@ -316,7 +317,7 @@ def _kernel( sfb_n_block = coord_n_pair // 128 if is_pair_leader: - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) for _ai in cutlass.range_constexpr(num_a_operands): @@ -327,7 +328,7 @@ def _kernel( sfa_m_block = coord_m_per_cta // 128 if cutlass.const_expr(multicast_a): if n_rank == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -354,7 +355,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_2, ) else: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -382,7 +383,7 @@ def _kernel( ) if cutlass.const_expr(multicast_a): if n_rank == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFA_stage, tma_sfa_desc.get_ptr(), @@ -393,7 +394,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_2, ) else: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFA_stage, tma_sfa_desc.get_ptr(), @@ -411,7 +412,7 @@ def _kernel( tma_sfb_desc = tma_sfb_descs[_bj] if cutlass.const_expr(multicast_b): if pair_m_idx == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -438,7 +439,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_2, ) else: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -466,7 +467,7 @@ def _kernel( ) if cutlass.const_expr(multicast_b): if pair_m_idx == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFB_stage, tma_sfb_desc.get_ptr(), @@ -477,7 +478,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_2, ) else: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFB_stage, tma_sfb_desc.get_ptr(), @@ -502,7 +503,7 @@ def _kernel( if tail_stage == ab_stages: tail_stage = cutlass.Int32(0) tail_phase = tail_phase ^ 1 - if nvvm.elect_sync(): + if elect_one: while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): pass @@ -638,7 +639,7 @@ def _kernel( for atom_r in cutlass.range(num_sf_atoms, unroll_full=True): for _ai in cutlass.range_constexpr(num_a_operands): for _m in cutlass.range_constexpr(num_blocks_m): - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_cp( s2t_shape, sfa_dst_ptrs[_ai][_m], @@ -648,7 +649,7 @@ def _kernel( ) for _bj in cutlass.range_constexpr(num_b_operands): for _m in cutlass.range_constexpr(num_blocks_n): - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_cp( s2t_shape, sfb_dst_ptrs[_bj][_m], @@ -664,7 +665,7 @@ def _kernel( _bj = gemm_b_idx[g] desc_a = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * k_block_idx) desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * k_block_idx) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_2, @@ -679,7 +680,7 @@ def _kernel( ) scale_d = cutlass.Boolean(True) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( ab_empty_mbar_ptr.subview(stage), multicast_mask=ab_empty_arrive_mask, @@ -687,7 +688,7 @@ def _kernel( ) ab_iter += 1 - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( acc_full_mbar_ptr.subview(acc_stage), multicast_mask=pair_mask, @@ -698,12 +699,12 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if nvvm.elect_sync(): + if elect_one: for _ in range(acc_stages): tail_stage = tail_stage + 1 if tail_stage == acc_stages: @@ -727,7 +728,7 @@ def _kernel( nvvm.tcgen05_dealloc(alloc_ptr, num_tmem_alloc_cols, group=nvvm.CTAGroup.CTA_2) else: if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) @@ -819,7 +820,7 @@ def _kernel( if use_acc_overlap and (not cd_out_is_m_major) and subtile_idx == acc_overlap_subtiles - 1: nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) @@ -876,7 +877,7 @@ def _kernel( ) if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(cd_out_is_m_major): for _mb in cutlass.range_constexpr(cta_tile_mnk[0] // cd_mmajor_atom_m): nvvm.cp_async_bulk_tensor_global_shared_cta( @@ -913,7 +914,7 @@ def _kernel( if cutlass.const_expr((not use_acc_overlap) or cd_out_is_m_major): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) @@ -923,7 +924,7 @@ def _kernel( if cutlass.const_expr(use_acc_overlap): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(tmem_dealloc_mbar_ptr) # @@TMA_STORE_ONLY:BEGIN@@ diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma.py index 7763af4000..8daf20b3ac 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma.py @@ -125,6 +125,7 @@ def _kernel( warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) + elect_one = nvvm.elect_sync() tidx = cute.arch.thread_idx()[0] bidx = cute.arch.block_idx()[0] @@ -240,7 +241,7 @@ def _kernel( num_consumer_warps_per_cta = 7 clc_empty_count = num_consumer_warps_per_cta * cluster_size if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: for i in range(ab_stages): nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) nvvm.mbarrier_init(ab_empty_mbar_ptr.subview(i), ab_empty_count) @@ -305,11 +306,11 @@ def _kernel( while not nvvm.mbarrier_try_wait_parity(clc_empty_mbar_ptr.subview(stage), clc_empty_phase, time_limit=10_000_000): pass - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(clc_full_mbar_ptr.subview(stage), 16) if is_cluster_leader_cta: - if nvvm.elect_sync(): + if elect_one: cute_clc.issue_clc_query( clc_full_mbar_cute_base + stage, clc_response_ptr_base + stage, @@ -324,7 +325,7 @@ def _kernel( is_valid_sched = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) @@ -347,7 +348,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) @@ -378,7 +379,7 @@ def _kernel( pass coord_k = k_tile_idx * cta_tile_mnk[2] - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) for _ai in cutlass.range_constexpr(num_a_operands): @@ -386,7 +387,7 @@ def _kernel( tma_a_desc = tma_a_descs[_ai] if cutlass.const_expr(multicast_a): if n_rank == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -413,7 +414,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_1, ) else: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -445,7 +446,7 @@ def _kernel( tma_b_desc = tma_b_descs[_bj] if cutlass.const_expr(multicast_b): if m_rank == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -472,7 +473,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_1, ) else: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -522,7 +523,7 @@ def _kernel( ) tile_l = l_idx nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 @@ -536,7 +537,7 @@ def _kernel( if tail_stage == ab_stages: tail_stage = cutlass.Int32(0) tail_phase = tail_phase ^ 1 - if nvvm.elect_sync(): + if elect_one: while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): pass @@ -602,7 +603,7 @@ def _kernel( stride_byte_offset=b_smem_desc_stride_byte_offset, layout=ab_smem_swizzle, ).advance_start_address(b_smem_k_step_bytes * k_block_idx) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_mma( mma_kind, nvvm.CTAGroup.CTA_1, @@ -614,7 +615,7 @@ def _kernel( ) scale_d = cutlass.Boolean(True) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( ab_empty_mbar_ptr.subview(stage), multicast_mask=ab_empty_arrive_mask, @@ -622,7 +623,7 @@ def _kernel( ) ab_iter += 1 - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( acc_full_mbar_ptr.subview(acc_stage), group=nvvm.CTAGroup.CTA_1, @@ -641,18 +642,18 @@ def _kernel( cute.arch.fence_proxy("async.shared", space="cta") is_valid = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if nvvm.elect_sync(): + if elect_one: for _ in range(acc_stages): tail_stage = tail_stage + 1 if tail_stage == acc_stages: @@ -739,7 +740,7 @@ def _kernel( if (not use_tma_store_epi) and subtile_idx == subtile_cnt - 1: nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) col = coord_n + subtile_col_offset @@ -798,7 +799,7 @@ def _kernel( ) if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(cd_out_is_m_major): for _mb in cutlass.range_constexpr(cta_tile_mnk[0] // cd_mmajor_atom_m): nvvm.cp_async_bulk_tensor_global_shared_cta( @@ -835,7 +836,7 @@ def _kernel( if cutlass.const_expr(use_tma_store_epi): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) consumer_stage = tile_iter % CLC_SCHED_STAGES @@ -859,7 +860,7 @@ def _kernel( ) tile_l = l_idx nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma_static.py b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma_static.py index cd9e75138e..714b05c0e4 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma_static.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma_static.py @@ -127,6 +127,7 @@ def _kernel( warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) + elect_one = nvvm.elect_sync() tidx = cute.arch.thread_idx()[0] bidx = cute.arch.block_idx()[0] @@ -226,7 +227,7 @@ def _kernel( ab_empty_count = cluster_m + cluster_n - 1 if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: for i in range(ab_stages): nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) nvvm.mbarrier_init(ab_empty_mbar_ptr.subview(i), ab_empty_count) @@ -277,7 +278,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) @@ -307,7 +308,7 @@ def _kernel( pass coord_k = k_tile_idx * cta_tile_mnk[2] - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) for _ai in cutlass.range_constexpr(num_a_operands): @@ -315,7 +316,7 @@ def _kernel( tma_a_desc = tma_a_descs[_ai] if cutlass.const_expr(multicast_a): if n_rank == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -342,7 +343,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_1, ) else: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -374,7 +375,7 @@ def _kernel( tma_b_desc = tma_b_descs[_bj] if cutlass.const_expr(multicast_b): if m_rank == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -401,7 +402,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_1, ) else: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -442,7 +443,7 @@ def _kernel( if tail_stage == ab_stages: tail_stage = cutlass.Int32(0) tail_phase = tail_phase ^ 1 - if nvvm.elect_sync(): + if elect_one: while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): pass @@ -507,7 +508,7 @@ def _kernel( stride_byte_offset=b_smem_desc_stride_byte_offset, layout=ab_smem_swizzle, ).advance_start_address(b_smem_k_step_bytes * k_block_idx) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_mma( mma_kind, nvvm.CTAGroup.CTA_1, @@ -519,7 +520,7 @@ def _kernel( ) scale_d = cutlass.Boolean(True) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( ab_empty_mbar_ptr.subview(stage), multicast_mask=ab_empty_arrive_mask, @@ -527,7 +528,7 @@ def _kernel( ) ab_iter += 1 - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( acc_full_mbar_ptr.subview(acc_stage), group=nvvm.CTAGroup.CTA_1, @@ -537,10 +538,10 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") - if nvvm.elect_sync(): + if elect_one: while not nvvm.mbarrier_try_wait_parity( acc_empty_mbar_ptr.subview(acc_stage), acc_empty_phase_bit ^ 1, @@ -624,7 +625,7 @@ def _kernel( if (not use_tma_store_epi) and subtile_idx == subtile_cnt - 1: nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) col = coord_n + subtile_col_offset @@ -683,7 +684,7 @@ def _kernel( ) if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(cd_out_is_m_major): for _mb in cutlass.range_constexpr(cta_tile_mnk[0] // cd_mmajor_atom_m): nvvm.cp_async_bulk_tensor_global_shared_cta( @@ -720,7 +721,7 @@ def _kernel( if cutlass.const_expr(use_tma_store_epi): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) is_valid = cutlass.Int32(0) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma.py index 259c8fc951..71ea815acd 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma.py @@ -128,6 +128,7 @@ def _kernel( warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) + elect_one = nvvm.elect_sync() tidx = cute.arch.thread_idx()[0] bidx = cute.arch.block_idx()[0] @@ -247,7 +248,7 @@ def _kernel( num_consumer_warps_per_cta = 7 clc_empty_count = num_consumer_warps_per_cta * cluster_size if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_init(tmem_dealloc_mbar_ptr, 32) for i in range(ab_stages): nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) @@ -315,11 +316,11 @@ def _kernel( while not nvvm.mbarrier_try_wait_parity(clc_empty_mbar_ptr.subview(stage), clc_empty_phase, time_limit=10_000_000): pass - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(clc_full_mbar_ptr.subview(stage), 16) if is_cluster_leader_cta: - if nvvm.elect_sync(): + if elect_one: cute_clc.issue_clc_query( clc_full_mbar_cute_base + stage, clc_response_ptr_base + stage, @@ -334,7 +335,7 @@ def _kernel( is_valid_sched = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) @@ -357,7 +358,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) @@ -391,14 +392,14 @@ def _kernel( coord_k = k_tile_idx * cgrp_tile_mnk[2] if is_pair_leader: - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) for _ai in cutlass.range_constexpr(num_a_operands): sA_stage = smem_a_list[_ai].subview(sA_elems * stage) tma_a_desc = tma_a_descs[_ai] if cutlass.const_expr(multicast_a): if n_rank == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -425,7 +426,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_2, ) else: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -457,7 +458,7 @@ def _kernel( tma_b_desc = tma_b_descs[_bj] if cutlass.const_expr(multicast_b): if pair_m_idx == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -484,7 +485,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_2, ) else: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -534,7 +535,7 @@ def _kernel( ) tile_l = l_idx nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 @@ -548,7 +549,7 @@ def _kernel( if tail_stage == ab_stages: tail_stage = cutlass.Int32(0) tail_phase = tail_phase ^ 1 - if nvvm.elect_sync(): + if elect_one: while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): pass @@ -631,7 +632,7 @@ def _kernel( stride_byte_offset=b_smem_desc_stride_byte_offset, layout=ab_smem_swizzle, ).advance_start_address(b_smem_k_step_bytes * k_block_idx) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_mma( mma_kind, nvvm.CTAGroup.CTA_2, @@ -643,7 +644,7 @@ def _kernel( ) scale_d = cutlass.Boolean(True) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( ab_empty_mbar_ptr.subview(stage), multicast_mask=ab_empty_arrive_mask, @@ -651,7 +652,7 @@ def _kernel( ) ab_iter += 1 - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( acc_full_mbar_ptr.subview(acc_stage), multicast_mask=pair_mask, @@ -671,18 +672,18 @@ def _kernel( cute.arch.fence_proxy("async.shared", space="cta") is_valid = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if nvvm.elect_sync(): + if elect_one: for _ in range(acc_stages): tail_stage = tail_stage + 1 if tail_stage == acc_stages: @@ -721,13 +722,13 @@ def _kernel( cute.arch.fence_proxy("async.shared", space="cta") is_valid = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) @@ -807,7 +808,7 @@ def _kernel( if (not use_tma_store_epi) and subtile_idx == subtile_cnt - 1: nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive( nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank), scope=nvvm.MemScope.CLUSTER, @@ -870,7 +871,7 @@ def _kernel( ) if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(cd_out_is_m_major): for _mb in cutlass.range_constexpr(cta_tile_mnk[0] // cd_mmajor_atom_m): nvvm.cp_async_bulk_tensor_global_shared_cta( @@ -907,7 +908,7 @@ def _kernel( if cutlass.const_expr(use_tma_store_epi): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) @@ -932,7 +933,7 @@ def _kernel( ) tile_l = l_idx nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma_static.py b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma_static.py index aa4a49cbf2..96d9cc1706 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma_static.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma_static.py @@ -125,6 +125,7 @@ def _kernel( warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) + elect_one = nvvm.elect_sync() tidx = cute.arch.thread_idx()[0] bidx = cute.arch.block_idx()[0] @@ -227,7 +228,7 @@ def _kernel( cta_group = 2 ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_init(tmem_dealloc_mbar_ptr, 32) for i in range(ab_stages): nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) @@ -281,7 +282,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) @@ -314,14 +315,14 @@ def _kernel( coord_k = k_tile_idx * cgrp_tile_mnk[2] if is_pair_leader: - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) for _ai in cutlass.range_constexpr(num_a_operands): sA_stage = smem_a_list[_ai].subview(sA_elems * stage) tma_a_desc = tma_a_descs[_ai] if cutlass.const_expr(multicast_a): if n_rank == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -348,7 +349,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_2, ) else: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -380,7 +381,7 @@ def _kernel( tma_b_desc = tma_b_descs[_bj] if cutlass.const_expr(multicast_b): if pair_m_idx == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -407,7 +408,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_2, ) else: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -448,7 +449,7 @@ def _kernel( if tail_stage == ab_stages: tail_stage = cutlass.Int32(0) tail_phase = tail_phase ^ 1 - if nvvm.elect_sync(): + if elect_one: while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): pass @@ -530,7 +531,7 @@ def _kernel( stride_byte_offset=b_smem_desc_stride_byte_offset, layout=ab_smem_swizzle, ).advance_start_address(b_smem_k_step_bytes * k_block_idx) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_mma( mma_kind, nvvm.CTAGroup.CTA_2, @@ -542,7 +543,7 @@ def _kernel( ) scale_d = cutlass.Boolean(True) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( ab_empty_mbar_ptr.subview(stage), multicast_mask=ab_empty_arrive_mask, @@ -550,7 +551,7 @@ def _kernel( ) ab_iter += 1 - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( acc_full_mbar_ptr.subview(acc_stage), multicast_mask=pair_mask, @@ -561,12 +562,12 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if nvvm.elect_sync(): + if elect_one: for _ in range(acc_stages): tail_stage = tail_stage + 1 if tail_stage == acc_stages: @@ -589,7 +590,7 @@ def _kernel( nvvm.tcgen05_dealloc(alloc_ptr, num_tmem_alloc_cols, group=nvvm.CTAGroup.CTA_2) else: if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) @@ -671,7 +672,7 @@ def _kernel( if (not use_tma_store_epi) and subtile_idx == subtile_cnt - 1: nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive( nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank), scope=nvvm.MemScope.CLUSTER, @@ -734,7 +735,7 @@ def _kernel( ) if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(cd_out_is_m_major): for _mb in cutlass.range_constexpr(cta_tile_mnk[0] // cd_mmajor_atom_m): nvvm.cp_async_bulk_tensor_global_shared_cta( @@ -771,7 +772,7 @@ def _kernel( if cutlass.const_expr(use_tma_store_epi): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_1ctamma.py index 8206e5146e..8e27be3386 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_1ctamma.py @@ -131,6 +131,7 @@ def _kernel( warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) + elect_one = nvvm.elect_sync() tidx = cute.arch.thread_idx()[0] bidx = cute.arch.block_idx()[0] @@ -244,7 +245,7 @@ def _kernel( num_consumer_warps_per_cta = 7 + num_mainloop_warps clc_empty_count = num_consumer_warps_per_cta * cluster_size if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: for i in range(ab_stages): nvvm.mbarrier_init(a_full_mbar_ptr.subview(i), 1) nvvm.mbarrier_init(b_full_mbar_ptr.subview(i), 1) @@ -319,11 +320,11 @@ def _kernel( while not nvvm.mbarrier_try_wait_parity(clc_empty_mbar_ptr.subview(stage), clc_empty_phase, time_limit=10_000_000): pass - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(clc_full_mbar_ptr.subview(stage), 16) if is_cluster_leader_cta: - if nvvm.elect_sync(): + if elect_one: cute_clc.issue_clc_query( clc_full_mbar_cute_base + stage, clc_response_ptr_base + stage, @@ -338,7 +339,7 @@ def _kernel( is_valid_sched = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) @@ -361,7 +362,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) @@ -402,13 +403,13 @@ def _kernel( sB_tma_dst = smem_b_load.subview(sB_elems * stage) else: sB_tma_dst = sB_stage - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(a_full_mbar_ptr.subview(stage), sA_tma_bytes) nvvm.mbarrier_arrive_expect_tx(b_full_mbar_ptr.subview(stage), sB_tma_bytes) if cutlass.const_expr(multicast_a): if n_rank == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -435,7 +436,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_1, ) else: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -464,7 +465,7 @@ def _kernel( if cutlass.const_expr(multicast_b): if m_rank == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -491,7 +492,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_1, ) else: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -541,7 +542,7 @@ def _kernel( ) tile_l = l_idx nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 @@ -555,7 +556,7 @@ def _kernel( if tail_stage == ab_stages: tail_stage = cutlass.Int32(0) tail_phase = tail_phase ^ 1 - if nvvm.elect_sync(): + if elect_one: while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): pass @@ -632,7 +633,7 @@ def _kernel( for k_block_idx in cutlass.range(num_k_blocks, unroll_full=True): desc_a = desc_a_base.advance_start_address(a_smem_k_step_bytes * k_block_idx) desc_b = desc_b_base.advance_start_address(b_smem_k_step_bytes * k_block_idx) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_mma( mma_kind, nvvm.CTAGroup.CTA_1, @@ -644,7 +645,7 @@ def _kernel( ) scale_d = cutlass.Boolean(True) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( ab_empty_mbar_ptr.subview(stage), multicast_mask=ab_empty_arrive_mask, @@ -652,7 +653,7 @@ def _kernel( ) ab_iter += 1 - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( acc_full_mbar_ptr.subview(acc_stage), group=nvvm.CTAGroup.CTA_1, @@ -671,18 +672,18 @@ def _kernel( cute.arch.fence_proxy("async.shared", space="cta") is_valid = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if nvvm.elect_sync(): + if elect_one: for _ in range(acc_stages): tail_stage = tail_stage + 1 if tail_stage == acc_stages: @@ -873,7 +874,7 @@ def _kernel( cute.arch.fence_view_async_shared() cute.arch.sync_warp() - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(mainloop_full_mbar_ptr.subview(stage)) ab_iter += 1 @@ -890,7 +891,7 @@ def _kernel( cute.arch.fence_proxy("async.shared", space="cta") is_valid = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 @@ -966,7 +967,7 @@ def _kernel( if (not use_tma_store_epi) and subtile_idx == subtile_cnt - 1: nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) col = coord_n + subtile_col_offset @@ -1025,7 +1026,7 @@ def _kernel( ) if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(cd_out_is_m_major): for _mb in cutlass.range_constexpr(cta_tile_mnk[0] // cd_mmajor_atom_m): nvvm.cp_async_bulk_tensor_global_shared_cta( @@ -1062,7 +1063,7 @@ def _kernel( if cutlass.const_expr(use_tma_store_epi): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) consumer_stage = tile_iter % CLC_SCHED_STAGES @@ -1086,7 +1087,7 @@ def _kernel( ) tile_l = l_idx nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_2ctamma.py index 56c3ecf5f4..97231eb2aa 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_2ctamma.py @@ -133,6 +133,7 @@ def _kernel( warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) + elect_one = nvvm.elect_sync() tidx = cute.arch.thread_idx()[0] bidx = cute.arch.block_idx()[0] @@ -250,7 +251,7 @@ def _kernel( clc_empty_count = num_consumer_warps_per_cta * cluster_size mainloop_full_count = num_mainloop_warps * 2 if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_init(tmem_dealloc_mbar_ptr, 32) for i in range(ab_stages): nvvm.mbarrier_init(a_full_mbar_ptr.subview(i), 1) @@ -328,11 +329,11 @@ def _kernel( while not nvvm.mbarrier_try_wait_parity(clc_empty_mbar_ptr.subview(stage), clc_empty_phase, time_limit=10_000_000): pass - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(clc_full_mbar_ptr.subview(stage), 16) if is_cluster_leader_cta: - if nvvm.elect_sync(): + if elect_one: cute_clc.issue_clc_query( clc_full_mbar_cute_base + stage, clc_response_ptr_base + stage, @@ -347,7 +348,7 @@ def _kernel( is_valid_sched = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) @@ -371,7 +372,7 @@ def _kernel( nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) @@ -415,22 +416,22 @@ def _kernel( sB_tma_dst = sB_stage if cutlass.const_expr(mainloop_fuse_a): - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(a_full_mbar_ptr.subview(stage), sA_tma_bytes) else: if is_pair_leader: - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(a_full_mbar_ptr.subview(stage), sA_tma_bytes * 2) if cutlass.const_expr(mainloop_fuse_b): - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(b_full_mbar_ptr.subview(stage), sB_tma_bytes) else: if is_pair_leader: - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(b_full_mbar_ptr.subview(stage), sB_tma_bytes * 2) if cutlass.const_expr(mainloop_fuse_a): a_self_mask = cutlass.Int16(1) << cta_rank_in_cluster - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -458,7 +459,7 @@ def _kernel( ) elif cutlass.const_expr(multicast_a): if n_rank == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -485,7 +486,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_2, ) else: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -514,7 +515,7 @@ def _kernel( if cutlass.const_expr(mainloop_fuse_b): b_self_mask = cutlass.Int16(1) << cta_rank_in_cluster - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -542,7 +543,7 @@ def _kernel( ) elif cutlass.const_expr(multicast_b): if pair_m_idx == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -569,7 +570,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_2, ) else: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -619,7 +620,7 @@ def _kernel( ) tile_l = l_idx nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 @@ -633,7 +634,7 @@ def _kernel( if tail_stage == ab_stages: tail_stage = cutlass.Int32(0) tail_phase = tail_phase ^ 1 - if nvvm.elect_sync(): + if elect_one: while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): pass @@ -731,7 +732,7 @@ def _kernel( for k_block_idx in cutlass.range(num_k_blocks, unroll_full=True): desc_a = desc_a_base.advance_start_address(a_smem_k_step_bytes * k_block_idx) desc_b = desc_b_base.advance_start_address(b_smem_k_step_bytes * k_block_idx) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_mma( mma_kind, nvvm.CTAGroup.CTA_2, @@ -743,7 +744,7 @@ def _kernel( ) scale_d = cutlass.Boolean(True) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( ab_empty_mbar_ptr.subview(stage), multicast_mask=ab_empty_arrive_mask, @@ -751,7 +752,7 @@ def _kernel( ) ab_iter += 1 - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( acc_full_mbar_ptr.subview(acc_stage), multicast_mask=pair_mask, @@ -771,18 +772,18 @@ def _kernel( cute.arch.fence_proxy("async.shared", space="cta") is_valid = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if nvvm.elect_sync(): + if elect_one: for _ in range(acc_stages): tail_stage = tail_stage + 1 if tail_stage == acc_stages: @@ -821,13 +822,13 @@ def _kernel( cute.arch.fence_proxy("async.shared", space="cta") is_valid = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) @@ -1010,7 +1011,7 @@ def _kernel( cute.arch.fence_view_async_shared() cute.arch.sync_warp() - if nvvm.elect_sync(): + if elect_one: ml_full_remote = nvvm.mapa(mainloop_full_mbar_ptr.subview(stage), pair_leader_rank) nvvm.mbarrier_arrive(ml_full_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) ab_iter += 1 @@ -1028,7 +1029,7 @@ def _kernel( cute.arch.fence_proxy("async.shared", space="cta") is_valid = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 @@ -1102,7 +1103,7 @@ def _kernel( if (not use_tma_store_epi) and subtile_idx == subtile_cnt - 1: nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive( nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank), scope=nvvm.MemScope.CLUSTER, @@ -1165,7 +1166,7 @@ def _kernel( ) if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(cd_out_is_m_major): for _mb in cutlass.range_constexpr(cta_tile_mnk[0] // cd_mmajor_atom_m): nvvm.cp_async_bulk_tensor_global_shared_cta( @@ -1202,7 +1203,7 @@ def _kernel( if cutlass.const_expr(use_tma_store_epi): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) @@ -1227,7 +1228,7 @@ def _kernel( ) tile_l = l_idx nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_1ctamma.py index 6de79b9ced..3b6b0b0619 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_1ctamma.py @@ -154,6 +154,7 @@ def _kernel( warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) + elect_one = nvvm.elect_sync() tidx = cute.arch.thread_idx()[0] bidx = cute.arch.block_idx()[0] @@ -268,7 +269,7 @@ def _kernel( ab_empty_count = cluster_m + cluster_n - 1 sched_empty_count = 1 + 1 + num_epilogue_warps if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: for i in range(ab_stages): nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) nvvm.mbarrier_init(ab_empty_mbar_ptr.subview(i), ab_empty_count) @@ -524,14 +525,13 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) sched_stage = cutlass.Int32(0) sched_full_phase = cutlass.Int32(0) is_valid = cutlass.Int32(1) - elect_one = nvvm.elect_sync() lane = tidx % 32 block_linear = bidx + bidy * gridx @@ -607,13 +607,13 @@ def _kernel( coord_k = k_tile_idx * cta_tile_mnk[2] coord_sf_k = k_tile_idx * sf_tma_box_k - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) a_issue = (not multicast_a) or (n_rank == 0) b_issue = (not multicast_b) or (m_rank == 0) if a_issue: for _ai in cutlass.range_constexpr(num_a_operands): - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( smem_a_list[_ai].subview(sA_elems * stage), a_desc_tma_ptr_list[_ai], @@ -624,7 +624,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_1, ) for _ai in cutlass.range_constexpr(num_a_operands): - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( smem_sfa_list[_ai].subview(sfa_smem_bytes * stage), tma_sfa_descs[_ai].get_ptr(), @@ -637,7 +637,7 @@ def _kernel( if b_issue: for _bj in cutlass.range_constexpr(num_b_operands): sB_stage = smem_b_list[_bj].subview(sB_elems * stage) - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -664,7 +664,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_1, ) for _bj in cutlass.range_constexpr(num_b_operands): - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( smem_sfb_list[_bj].subview(sfb_smem_bytes * stage), tma_sfb_descs[_bj].get_ptr(), @@ -685,7 +685,7 @@ def _kernel( if tail_stage == ab_stages: tail_stage = cutlass.Int32(0) tail_phase = tail_phase ^ 1 - if nvvm.elect_sync(): + if elect_one: while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): pass @@ -740,7 +740,7 @@ def _kernel( ): pass is_valid = (sched_storage.subview(sched_stage * SCHED_SLOT_WORDS).subview(3)).load() - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(sched_empty_mbar_ptr.subview(sched_stage)) sched_stage += 1 if sched_stage == SCHED_STAGES: @@ -824,7 +824,7 @@ def _kernel( for atom_r in cutlass.range_constexpr(num_sf_atoms): for _ai in cutlass.range_constexpr(num_a_operands): for _m in cutlass.range_constexpr(num_blocks_m): - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_cp( s2t_shape, sfa_dst_ptrs[_ai][_m], @@ -834,7 +834,7 @@ def _kernel( ) for _bj in cutlass.range_constexpr(num_b_operands): for _m in cutlass.range_constexpr(num_blocks_n): - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_cp( s2t_shape, sfb_dst_ptrs[_bj][_m], @@ -850,7 +850,7 @@ def _kernel( _bj = gemm_b_idx[g] desc_a = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * k_block_idx) desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * k_block_idx) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_1, @@ -865,7 +865,7 @@ def _kernel( ) scale_d = cutlass.Boolean(True) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( ab_empty_mbar_ptr.subview(stage), multicast_mask=ab_empty_arrive_mask, @@ -873,7 +873,7 @@ def _kernel( ) ab_iter += 1 - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( acc_full_mbar_ptr.subview(acc_stage), group=nvvm.CTAGroup.CTA_1, @@ -881,14 +881,14 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) if tile_iter != 0: tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if nvvm.elect_sync(): + if elect_one: for _ in range(acc_stages): tail_stage = tail_stage + 1 if tail_stage == acc_stages: @@ -948,7 +948,7 @@ def _kernel( group_begin = (_slot.subview(4)).load() group_end = (_slot.subview(5)).load() group_idx = (_slot.subview(7)).load() - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(sched_empty_mbar_ptr.subview(sched_stage)) sched_stage += 1 if sched_stage == SCHED_STAGES: @@ -1005,7 +1005,7 @@ def _kernel( if cutlass.const_expr(use_acc_overlap and subtile_idx == acc_overlap_subtiles - 1): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) col = coord_n + subtile_col_offset @@ -1024,13 +1024,13 @@ def _kernel( if cutlass.const_expr(not use_acc_overlap): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) tile_iter += 1 if cutlass.const_expr(use_acc_overlap): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(tmem_dealloc_mbar_ptr) if warp_idx == unused_warp_id: diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_2ctamma.py index 2c49dd2368..1d739d3882 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_2ctamma.py @@ -156,6 +156,7 @@ def _kernel( warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) + elect_one = nvvm.elect_sync() tidx = cute.arch.thread_idx()[0] bidx = cute.arch.block_idx()[0] @@ -273,7 +274,7 @@ def _kernel( ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 sched_empty_count = 1 + 1 + num_epilogue_warps if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(use_acc_overlap): nvvm.mbarrier_init(tmem_dealloc_mbar_ptr, num_epilogue_warps) else: @@ -532,14 +533,13 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) sched_stage = cutlass.Int32(0) sched_full_phase = cutlass.Int32(0) is_valid = cutlass.Int32(1) - elect_one = nvvm.elect_sync() logical_cta_tile_n = cgrp_tile_mnk[1] // cluster_n lane = tidx % 32 @@ -619,13 +619,13 @@ def _kernel( coord_sf_k = k_tile_idx * sf_tma_box_k if is_pair_leader: - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) a_issue = (not multicast_a) or (n_rank == 0) b_issue = (not multicast_b) or (pair_m_idx == 0) if a_issue: for _ai in cutlass.range_constexpr(num_a_operands): - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( smem_a_list[_ai].subview(sA_elems * stage), a_desc_tma_ptr_list[_ai], @@ -636,7 +636,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_2, ) for _ai in cutlass.range_constexpr(num_a_operands): - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( smem_sfa_list[_ai].subview(sfa_smem_bytes * stage), tma_sfa_descs[_ai].get_ptr(), @@ -649,7 +649,7 @@ def _kernel( if b_issue: for _bj in cutlass.range_constexpr(num_b_operands): sB_stage = smem_b_list[_bj].subview(sB_elems * stage) - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -676,7 +676,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_2, ) for _bj in cutlass.range_constexpr(num_b_operands): - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( smem_sfb_list[_bj].subview(sfb_smem_bytes * stage), tma_sfb_descs[_bj].get_ptr(), @@ -698,7 +698,7 @@ def _kernel( if tail_stage == ab_stages: tail_stage = cutlass.Int32(0) tail_phase = tail_phase ^ 1 - if nvvm.elect_sync(): + if elect_one: while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): pass @@ -764,7 +764,7 @@ def _kernel( ): pass is_valid = (sched_storage.subview(sched_stage * SCHED_SLOT_WORDS).subview(3)).load() - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(sched_empty_mbar_ptr.subview(sched_stage)) sched_stage += 1 if sched_stage == SCHED_STAGES: @@ -848,7 +848,7 @@ def _kernel( for atom_r in cutlass.range(num_sf_atoms, unroll_full=True): for _ai in cutlass.range_constexpr(num_a_operands): for _m in cutlass.range_constexpr(num_blocks_m): - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_cp( s2t_shape, sfa_dst_ptrs[_ai][_m], @@ -858,7 +858,7 @@ def _kernel( ) for _bj in cutlass.range_constexpr(num_b_operands): for _m in cutlass.range_constexpr(num_blocks_n): - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_cp( s2t_shape, sfb_dst_ptrs[_bj][_m], @@ -874,7 +874,7 @@ def _kernel( _bj = gemm_b_idx[g] desc_a = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * k_block_idx) desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * k_block_idx) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_2, @@ -889,7 +889,7 @@ def _kernel( ) scale_d = cutlass.Boolean(True) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( ab_empty_mbar_ptr.subview(stage), multicast_mask=ab_empty_arrive_mask, @@ -897,7 +897,7 @@ def _kernel( ) ab_iter += 1 - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( acc_full_mbar_ptr.subview(acc_stage), multicast_mask=pair_mask, @@ -906,13 +906,13 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") if tile_iter != 0: tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if nvvm.elect_sync(): + if elect_one: for _ in range(acc_stages): tail_stage = tail_stage + 1 if tail_stage == acc_stages: @@ -946,7 +946,7 @@ def _kernel( ): pass is_valid = (sched_storage.subview(sched_stage * SCHED_SLOT_WORDS).subview(3)).load() - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(sched_empty_mbar_ptr.subview(sched_stage)) sched_stage += 1 if sched_stage == SCHED_STAGES: @@ -954,7 +954,7 @@ def _kernel( sched_full_phase = sched_full_phase ^ 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) @@ -1004,7 +1004,7 @@ def _kernel( group_begin = (slot.subview(4)).load() group_end = (slot.subview(5)).load() group_idx = (slot.subview(7)).load() - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(sched_empty_mbar_ptr.subview(sched_stage)) sched_stage += 1 if sched_stage == SCHED_STAGES: @@ -1063,7 +1063,7 @@ def _kernel( if use_acc_overlap and subtile_idx == acc_overlap_subtiles - 1: nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) @@ -1085,14 +1085,14 @@ def _kernel( if cutlass.const_expr(not use_acc_overlap): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 if cutlass.const_expr(use_acc_overlap): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(tmem_dealloc_mbar_ptr) if warp_idx == unused_warp_id: diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_1ctamma.py index ff85fd973e..0ee86590af 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_1ctamma.py @@ -189,6 +189,7 @@ def _kernel( warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) + elect_one = nvvm.elect_sync() tidx = cute.arch.thread_idx()[0] bidx = cute.arch.block_idx()[0] @@ -282,7 +283,7 @@ def _kernel( ab_empty_count = cluster_m + cluster_n - 1 sched_empty_count = 1 + 1 + num_epilogue_warps if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: for i in range(ab_stages): nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) nvvm.mbarrier_init(ab_empty_mbar_ptr.subview(i), ab_empty_count) @@ -443,14 +444,13 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) sched_stage = cutlass.Int32(0) sched_full_phase = cutlass.Int32(0) is_valid = cutlass.Int32(1) - elect_one = nvvm.elect_sync() lane = tidx % 32 block_linear = bidx + bidy * gridx @@ -522,14 +522,14 @@ def _kernel( pass coord_k = k_tile_idx * cta_tile_mnk[2] - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) a_issue = (not multicast_a) or (n_rank == 0) if a_issue: for _ai in cutlass.range_constexpr(num_a_operands): sA_stage = smem_a_list[_ai].subview(sA_elems * stage) - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sA_stage, a_desc_tma_ptr_list[_ai], @@ -543,7 +543,7 @@ def _kernel( if b_issue: for _bj in cutlass.range_constexpr(num_b_operands): sB_stage = smem_b_list[_bj].subview(sB_elems * stage) - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -580,7 +580,7 @@ def _kernel( if tail_stage == ab_stages: tail_stage = cutlass.Int32(0) tail_phase = tail_phase ^ 1 - if nvvm.elect_sync(): + if elect_one: while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): pass @@ -608,7 +608,7 @@ def _kernel( ): pass is_valid = (sched_storage.subview(sched_stage * SCHED_SLOT_WORDS).subview(3)).load() - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(sched_empty_mbar_ptr.subview(sched_stage)) sched_stage += 1 if sched_stage == SCHED_STAGES: @@ -666,7 +666,7 @@ def _kernel( stride_byte_offset=b_smem_desc_stride_byte_offset, layout=ab_smem_swizzle, ).advance_start_address(b_smem_k_step_bytes * k_block_idx) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_mma( mma_kind, nvvm.CTAGroup.CTA_1, @@ -678,7 +678,7 @@ def _kernel( ) scale_d = cutlass.Boolean(True) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( ab_empty_mbar_ptr.subview(stage), multicast_mask=ab_empty_arrive_mask, @@ -686,7 +686,7 @@ def _kernel( ) ab_iter += 1 - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( acc_full_mbar_ptr.subview(acc_stage), group=nvvm.CTAGroup.CTA_1, @@ -694,13 +694,13 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") if tile_iter != 0: tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if nvvm.elect_sync(): + if elect_one: for _ in range(acc_stages): tail_stage = tail_stage + 1 if tail_stage == acc_stages: @@ -751,7 +751,7 @@ def _kernel( group_begin = (_slot.subview(4)).load() group_end = (_slot.subview(5)).load() group_idx = (_slot.subview(7)).load() - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(sched_empty_mbar_ptr.subview(sched_stage)) sched_stage += 1 if sched_stage == SCHED_STAGES: @@ -796,7 +796,7 @@ def _kernel( if subtile_idx == subtile_cnt - 1: nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) col = coord_n + subtile_col_offset @@ -828,7 +828,7 @@ def _kernel( group_begin = (_slot.subview(4)).load() group_end = (_slot.subview(5)).load() group_idx = (_slot.subview(7)).load() - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(sched_empty_mbar_ptr.subview(sched_stage)) sched_stage += 1 if sched_stage == SCHED_STAGES: diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_2ctamma.py index 97ff5b7239..14055d0d28 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_2ctamma.py @@ -194,6 +194,7 @@ def _kernel( warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) + elect_one = nvvm.elect_sync() tidx = cute.arch.thread_idx()[0] bidx = cute.arch.block_idx()[0] @@ -293,7 +294,7 @@ def _kernel( num_consumer_warps_per_cta = 1 + 1 + num_epilogue_warps sched_empty_count = num_consumer_warps_per_cta if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_init(tmem_dealloc_mbar_ptr, 32) for i in range(ab_stages): nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) @@ -458,14 +459,13 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) sched_stage = cutlass.Int32(0) sched_full_phase = cutlass.Int32(0) is_valid = cutlass.Int32(1) - elect_one = nvvm.elect_sync() logical_cta_tile_n = cgrp_tile_mnk[1] // cluster_n lane = tidx % 32 @@ -540,13 +540,13 @@ def _kernel( coord_k = k_tile_idx * cgrp_tile_mnk[2] if is_pair_leader: - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) a_issue = (not multicast_a) or (n_rank == 0) if a_issue: for _ai in cutlass.range_constexpr(num_a_operands): sA_stage = smem_a_list[_ai].subview(sA_elems * stage) - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sA_stage, a_desc_tma_ptr_list[_ai], @@ -560,7 +560,7 @@ def _kernel( if b_issue: for _bj in cutlass.range_constexpr(num_b_operands): sB_stage = smem_b_list[_bj].subview(sB_elems * stage) - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -597,7 +597,7 @@ def _kernel( if tail_stage == ab_stages: tail_stage = cutlass.Int32(0) tail_phase = tail_phase ^ 1 - if nvvm.elect_sync(): + if elect_one: while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): pass @@ -638,7 +638,7 @@ def _kernel( ): pass is_valid = (sched_storage.subview(sched_stage * SCHED_SLOT_WORDS).subview(3)).load() - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(sched_empty_mbar_ptr.subview(sched_stage)) sched_stage += 1 if sched_stage == SCHED_STAGES: @@ -696,7 +696,7 @@ def _kernel( stride_byte_offset=b_smem_desc_stride_byte_offset, layout=ab_smem_swizzle, ).advance_start_address(b_smem_k_step_bytes * k_block_idx) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_mma( mma_kind, nvvm.CTAGroup.CTA_2, @@ -708,7 +708,7 @@ def _kernel( ) scale_d = cutlass.Boolean(True) - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( ab_empty_mbar_ptr.subview(stage), multicast_mask=ab_empty_arrive_mask, @@ -716,7 +716,7 @@ def _kernel( ) ab_iter += 1 - if nvvm.elect_sync(): + if elect_one: nvvm.tcgen05_commit( acc_full_mbar_ptr.subview(acc_stage), multicast_mask=pair_mask, @@ -725,13 +725,13 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") if tile_iter != 0: tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if nvvm.elect_sync(): + if elect_one: for _ in range(acc_stages): tail_stage = tail_stage + 1 if tail_stage == acc_stages: @@ -764,7 +764,7 @@ def _kernel( ): pass is_valid = (sched_storage.subview(sched_stage * SCHED_SLOT_WORDS).subview(3)).load() - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(sched_empty_mbar_ptr.subview(sched_stage)) sched_stage += 1 if sched_stage == SCHED_STAGES: @@ -772,7 +772,7 @@ def _kernel( sched_full_phase = sched_full_phase ^ 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) @@ -817,7 +817,7 @@ def _kernel( group_begin = (slot.subview(4)).load() group_end = (slot.subview(5)).load() group_idx = (slot.subview(7)).load() - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(sched_empty_mbar_ptr.subview(sched_stage)) sched_stage += 1 if sched_stage == SCHED_STAGES: @@ -868,7 +868,7 @@ def _kernel( if subtile_idx == subtile_cnt - 1: nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive( nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank), scope=nvvm.MemScope.CLUSTER, diff --git a/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_1ctamma.py index be62539b5b..fa334d391b 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_1ctamma.py @@ -180,6 +180,7 @@ def _kernel( warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) + elect_one = nvvm.elect_sync() tidx = cute.arch.thread_idx()[0] bidx = cute.arch.block_idx()[0] @@ -318,7 +319,7 @@ def _kernel( num_consumer_warps_per_cta = 8 clc_empty_count = num_consumer_warps_per_cta * cluster_size if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: for i in range(ab_stages): nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) nvvm.mbarrier_init(ab_empty_mbar_ptr.subview(i), ab_empty_count) @@ -375,11 +376,11 @@ def _kernel( while not nvvm.mbarrier_try_wait_parity(clc_empty_mbar_ptr.subview(stage), clc_empty_phase, time_limit=10_000_000): pass - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(clc_full_mbar_ptr.subview(stage), 16) if is_cluster_leader_cta: - if nvvm.elect_sync(): + if elect_one: cute_clc.issue_clc_query( clc_full_mbar_cute_base + stage, clc_response_ptr_base + stage, @@ -394,7 +395,7 @@ def _kernel( is_valid_sched = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) @@ -417,7 +418,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_stage_cur = cutlass.Int32(0) # incremental ring walk — no div in the loop @@ -447,7 +448,7 @@ def _kernel( pass coord_k = k_tile_idx * cta_tile_mnk[2] + _kc * ab_tma_box_k_elems - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_ab_chunk_bytes) for _ai in cutlass.range_constexpr(num_a_operands): @@ -455,7 +456,7 @@ def _kernel( tma_a_desc = tma_a_descs[_ai] if cutlass.const_expr(multicast_a): if n_rank == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sA_stage, tma_a_desc.get_ptr(), @@ -466,7 +467,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_1, ) else: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sA_stage, tma_a_desc.get_ptr(), @@ -482,7 +483,7 @@ def _kernel( tma_b_desc = tma_b_descs[_bj] if cutlass.const_expr(multicast_b): if m_rank == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sB_stage, tma_b_desc.get_ptr(), @@ -493,7 +494,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_1, ) else: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sB_stage, tma_b_desc.get_ptr(), @@ -529,7 +530,7 @@ def _kernel( ) tile_l = l_idx nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 @@ -543,14 +544,14 @@ def _kernel( if tail_stage == ab_stages: tail_stage = cutlass.Int32(0) tail_phase = tail_phase ^ 1 - if nvvm.elect_sync(): + if elect_one: while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): pass if warp_idx == sf_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("wait") sf_empty_phase_bit = cutlass.Int32(1) sf_stage_cur = cutlass.Int32(0) # incremental ring walk — no div in the loop @@ -582,7 +583,7 @@ def _kernel( pass coord_sf_k = k_tile_idx * sf_tma_box_k + _grp * sf_atoms_per_group - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(sf_full_mbar_ptr.subview(stage), num_tma_sf_group_bytes) for _ai in cutlass.range_constexpr(num_a_operands): @@ -590,7 +591,7 @@ def _kernel( tma_sfa_desc = tma_sfa_descs[_ai] if cutlass.const_expr(multicast_a): if n_rank == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFA_stage, tma_sfa_desc.get_ptr(), @@ -601,7 +602,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_1, ) else: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFA_stage, tma_sfa_desc.get_ptr(), @@ -617,7 +618,7 @@ def _kernel( tma_sfb_desc = tma_sfb_descs[_bj] if cutlass.const_expr(multicast_b): if m_rank == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFB_stage, tma_sfb_desc.get_ptr(), @@ -628,7 +629,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_1, ) else: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFB_stage, tma_sfb_desc.get_ptr(), @@ -664,7 +665,7 @@ def _kernel( ) tile_l = l_idx nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 @@ -676,7 +677,7 @@ def _kernel( if tail_stage == sf_stages: tail_stage = cutlass.Int32(0) tail_phase = tail_phase ^ 1 - if nvvm.elect_sync(): + if elect_one: while not nvvm.mbarrier_try_wait_parity(sf_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): pass @@ -741,7 +742,7 @@ def _kernel( ): pass - is_mma_leader = nvvm.elect_sync() + is_mma_leader = elect_one if cutlass.const_expr(use_acc_overlap): acc_base_col = base_col_id_root + (tile_iter % 2) * acc_stage_stride else: @@ -963,18 +964,18 @@ def _kernel( cute.arch.fence_proxy("async.shared", space="cta") is_valid = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if nvvm.elect_sync(): + if elect_one: for _ in range(acc_stages): tail_stage = tail_stage + 1 if tail_stage == acc_stages: @@ -1067,7 +1068,7 @@ def _kernel( if use_acc_overlap and (not cd_out_is_m_major) and subtile_idx == acc_overlap_subtiles - 1: nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) col = coord_n + subtile_col_offset @@ -1123,7 +1124,7 @@ def _kernel( ) if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(cd_out_is_m_major): for _mb in cutlass.range_constexpr(cta_tile_mnk[0] // cd_mmajor_atom_m): nvvm.cp_async_bulk_tensor_global_shared_cta( @@ -1160,7 +1161,7 @@ def _kernel( if cutlass.const_expr((not use_acc_overlap) or cd_out_is_m_major): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) consumer_stage = tile_iter % CLC_SCHED_STAGES @@ -1184,7 +1185,7 @@ def _kernel( ) tile_l = l_idx nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) @@ -1192,7 +1193,7 @@ def _kernel( if cutlass.const_expr(use_acc_overlap): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(tmem_dealloc_mbar_ptr) # @@TMA_STORE_ONLY:BEGIN@@ diff --git a/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_2ctamma.py index 38de8bdb55..4fd1228dd1 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_2ctamma.py @@ -180,6 +180,7 @@ def _kernel( warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) + elect_one = nvvm.elect_sync() tidx = cute.arch.thread_idx()[0] bidx = cute.arch.block_idx()[0] @@ -320,7 +321,7 @@ def _kernel( num_consumer_warps_per_cta = 8 clc_empty_count = num_consumer_warps_per_cta * cluster_size if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(use_acc_overlap): nvvm.mbarrier_init(tmem_dealloc_mbar_ptr, num_epilogue_warps) else: @@ -384,11 +385,11 @@ def _kernel( while not nvvm.mbarrier_try_wait_parity(clc_empty_mbar_ptr.subview(stage), clc_empty_phase, time_limit=10_000_000): pass - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(clc_full_mbar_ptr.subview(stage), 16) if is_cluster_leader_cta: - if nvvm.elect_sync(): + if elect_one: cute_clc.issue_clc_query( clc_full_mbar_cute_base + stage, clc_response_ptr_base + stage, @@ -403,7 +404,7 @@ def _kernel( is_valid_sched = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) @@ -426,7 +427,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_stage_cur = cutlass.Int32(0) # incremental ring walk — no div in the loop @@ -458,7 +459,7 @@ def _kernel( coord_k = k_tile_idx * cta_tile_mnk[2] + _kc * ab_tma_box_k_elems if is_pair_leader: - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_ab_chunk_bytes) for _ai in cutlass.range_constexpr(num_a_operands): @@ -466,7 +467,7 @@ def _kernel( tma_a_desc = tma_a_descs[_ai] if cutlass.const_expr(multicast_a): if n_rank == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sA_stage, tma_a_desc.get_ptr(), @@ -477,7 +478,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_2, ) else: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sA_stage, tma_a_desc.get_ptr(), @@ -493,7 +494,7 @@ def _kernel( tma_b_desc = tma_b_descs[_bj] if cutlass.const_expr(multicast_b): if pair_m_idx == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sB_stage, tma_b_desc.get_ptr(), @@ -504,7 +505,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_2, ) else: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sB_stage, tma_b_desc.get_ptr(), @@ -540,7 +541,7 @@ def _kernel( ) tile_l = l_idx nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 @@ -554,14 +555,14 @@ def _kernel( if tail_stage == ab_stages: tail_stage = cutlass.Int32(0) tail_phase = tail_phase ^ 1 - if nvvm.elect_sync(): + if elect_one: while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): pass if warp_idx == sf_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("wait") sf_empty_phase_bit = cutlass.Int32(1) sf_stage_cur = cutlass.Int32(0) # incremental ring walk — no div in the loop @@ -596,7 +597,7 @@ def _kernel( coord_sf_k = k_tile_idx * sf_tma_box_k + _grp * sf_atoms_per_group if is_pair_leader: - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive_expect_tx(sf_full_mbar_ptr.subview(stage), num_tma_sf_group_bytes) for _ai in cutlass.range_constexpr(num_a_operands): @@ -604,7 +605,7 @@ def _kernel( tma_sfa_desc = tma_sfa_descs[_ai] if cutlass.const_expr(multicast_a): if n_rank == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFA_stage, tma_sfa_desc.get_ptr(), @@ -615,7 +616,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_2, ) else: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFA_stage, tma_sfa_desc.get_ptr(), @@ -631,7 +632,7 @@ def _kernel( tma_sfb_desc = tma_sfb_descs[_bj] if cutlass.const_expr(multicast_b): if pair_m_idx == 0: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFB_stage, tma_sfb_desc.get_ptr(), @@ -642,7 +643,7 @@ def _kernel( group=nvvm.CTAGroup.CTA_2, ) else: - if nvvm.elect_sync(): + if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( sSFB_stage, tma_sfb_desc.get_ptr(), @@ -678,7 +679,7 @@ def _kernel( ) tile_l = l_idx nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 @@ -690,7 +691,7 @@ def _kernel( if tail_stage == sf_stages: tail_stage = cutlass.Int32(0) tail_phase = tail_phase ^ 1 - if nvvm.elect_sync(): + if elect_one: while not nvvm.mbarrier_try_wait_parity(sf_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): pass @@ -768,7 +769,7 @@ def _kernel( ): pass - is_mma_leader = nvvm.elect_sync() + is_mma_leader = elect_one if cutlass.const_expr(use_acc_overlap): acc_base_col = base_col_id_root + (tile_iter % 2) * acc_stage_stride else: @@ -991,18 +992,18 @@ def _kernel( cute.arch.fence_proxy("async.shared", space="cta") is_valid = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if nvvm.elect_sync(): + if elect_one: for _ in range(acc_stages): tail_stage = tail_stage + 1 if tail_stage == acc_stages: @@ -1041,13 +1042,13 @@ def _kernel( cute.arch.fence_proxy("async.shared", space="cta") is_valid = vld nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) tile_iter += 1 if cutlass.const_expr(USE_PDL): - if nvvm.elect_sync(): + if elect_one: nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) @@ -1138,7 +1139,7 @@ def _kernel( if use_acc_overlap and (not cd_out_is_m_major) and subtile_idx == acc_overlap_subtiles - 1: nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) @@ -1195,7 +1196,7 @@ def _kernel( ) if warp_idx == 0: - if nvvm.elect_sync(): + if elect_one: if cutlass.const_expr(cd_out_is_m_major): for _mb in cutlass.range_constexpr(cta_tile_mnk[0] // cd_mmajor_atom_m): nvvm.cp_async_bulk_tensor_global_shared_cta( @@ -1232,7 +1233,7 @@ def _kernel( if cutlass.const_expr((not use_acc_overlap) or cd_out_is_m_major): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) @@ -1257,7 +1258,7 @@ def _kernel( ) tile_l = l_idx nvvm.bar_warp_sync(0xFFFFFFFF) - if nvvm.elect_sync(): + if elect_one: empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0) nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True) @@ -1265,7 +1266,7 @@ def _kernel( if cutlass.const_expr(use_acc_overlap): nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if nvvm.elect_sync(): + if elect_one: nvvm.mbarrier_arrive(tmem_dealloc_mbar_ptr) # @@TMA_STORE_ONLY:BEGIN@@